BZOJ 1468: Tree

题目大意：

问树上有多少对点他们的距离小于等于k。

题解：
点分治。

代码：

#include<cstdio>
#include<algorithm>
using namespace std;
long long ans;
int n,k,cnt,root,num,tail,last[1000005],sz[1000005],f[1000005],q[1000005],dis[1000005],vis[1000005];
struct node{
	int to,next,val;
}e[1000005];
void add(int a,int b,int c){
	e[++cnt].to=b;
	e[cnt].next=last[a];
	e[cnt].val=c;
	last[a]=cnt;
}
void getroot(int x,int fa){
	sz[x]=1;
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (V==fa || vis[V]) continue;
		getroot(V,x);
		sz[x]+=sz[V];
		f[x]=max(f[x],sz[V]);
	}
	f[x]=max(f[x],num-f[x]);
	if (f[root]>f[x]) root=x;
}
void getdis(int x,int fa){
	q[++tail]=dis[x];
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (V==fa || vis[V]) continue;
		dis[V]=dis[x]+e[i].val;
		getdis(V,x);
	}
}
long long calc(int x,int val){
	tail=0;
	dis[x]=val;
	getdis(x,0);
	long long sum=0;
	int head=1;
	sort(q+1,q+tail+1);
	while (head<tail){
		if (q[head]+q[tail]<=k) {
			sum+=tail-head;
			head++;
		}
		else tail--;
	}
	return sum;
}
void dfs(int x){
	ans+=calc(x,0);
	vis[x]=1;
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (vis[V]) continue;
		ans-=calc(V,e[i].val);
	}
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (vis[V]) continue;
		root=0;
		num=sz[V];
		getroot(V,0);
		dfs(root);
	}
}
int main(){
	scanf("%d",&n);
	for (int i=1; i<n; i++){
		int a,b,c;
		scanf("%d%d%d",&a,&b,&c);
		add(a,b,c);
		add(b,a,c);
	}
	scanf("%d",&k);
	num=n;
	root=0;
	f[root]=1e9; 
	getroot(1,0);
	dfs(root);
	printf("%lld\n",ans);
	return 0;
}