luogu P4178 Tree |点分治+树状数组

题目描述

给你一棵TREE,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K
输入格式

N（n<=40000） 接下来n-1行边描述管道，按照题目中写的输入 接下来是k

输出格式

一行，有多少对点之间的距离小于等于k

---

如此简洁的题目在洛谷是真的少见啊

点分治统计，树状数组快速查找小于K的 O(nlogn*logn)

#include<cstdio>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
const int N=5e5+10,inf=4e7;
using namespace std;
int n,K;
int nxt[N<<1],head[N],go[N<<1],w[N<<1],tot;
inline void add(int u,int v,int o){
	nxt[++tot]=head[u]; head[u]=tot; go[tot]=v; w[tot]=o;
	nxt[++tot]=head[v]; head[v]=tot; go[tot]=u; w[tot]=o;
}
int maxp[N],size[N],dis[N],rem[N];
int vis[N];
int sum,rt;
inline void getrt(int u,int fa){
	size[u]=1; maxp[u]=0;
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(v==fa)continue;
		size[u]+=size[v];
		maxp[u]=max(maxp[u],size[v]);
	}
	maxp[u]=max(maxp[u],sum-size[u]);
	if(maxp[u]<maxp[rt])rt=u;
}
inline void getdis(int u,int fa){
	if(dis[u]>K)return;
	rem[++rem[0]]=dis[u];
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(v==fa||vis[v])continue;
		dis[v]=dis[u]+w[i];
		getdis(v,u);
	}
}
int c[N];
inline void update(int x,int y){
	for(;x<=K;x+=x&(-x))c[x]+=y;
}
inline int ask(int x){
	int ans=0;
	for(;x;x-=x&(-x))ans+=c[x];
	return ans;
}
int ans=0,q[N];
inline void calc(int u){
	int p=0;
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(vis[v])continue;
		rem[0]=0; dis[v]=w[i];
		getdis(v,u);
		for(int j=rem[0];j;j--)ans+=ask(K-rem[j])+1;
		for(int j=rem[0];j;j--)
		update(rem[j],1),q[++p]=rem[j];
	}
	for(int i=1;i<=p;i++)update(q[i],ask(q[i]-1)-ask(q[i]));
}
inline void solve(int u){
	vis[u]=1; calc(u);
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(vis[v])continue;
		sum=size[v];
		maxp[rt=0]=inf;
		getrt(v,0); 
		solve(rt);
	}
}
signed main(){
	cin>>n;
	for(int i=1,u,v,o;i<n;i++){
		scanf("%d%d%d",&u,&v,&o);
		add(u,v,o);
	}
	cin>>K; 
	maxp[0]=sum=n;
	getrt(1,0);
	solve(rt);
	cout<<ans<<endl;
}