poj1741 Tree

参考博文
带一点容斥的思想

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
using namespace std;
struct Edge{
	int too, nxt, val;
}edge[20005];
int n, k, uu, vv, ww, cnt, sze, rot, hea[10005], gra[10005], siz[10005];
int ans, tot, dep[10005], d[10005];
bool vis[10005];
void add_edge(int fro, int too, int val){
	edge[++cnt].nxt = hea[fro];
	edge[cnt].too = too;
	edge[cnt].val = val;
	hea[fro] = cnt;
}
void init(){
	memset(hea, 0, sizeof(hea));
	memset(vis, 0, sizeof(vis));
	ans = cnt = sze = rot = 0;
}
void getRoot(int x, int fa){
	gra[x] = 0;
	siz[x] = 1;
	for(int i=hea[x]; i; i=edge[i].nxt){
		int t=edge[i].too;
		if(t!=fa && !vis[t]){
			getRoot(t, x);
			siz[x] += siz[t];
			gra[x] = max(gra[x], siz[t]);
		}
	}
	gra[x] = max(gra[x], sze-siz[x]);
	if(gra[rot]>gra[x])	rot = x;
}
void getDeep(int x, int fa){
	d[++tot] = dep[x];
	for(int i=hea[x]; i; i=edge[i].nxt){
		int t=edge[i].too;
		if(t!=fa && !vis[t]){
			dep[t] = dep[x] + edge[i].val;
			getDeep(t, x);
		}
	}
}
int calc(int x){
	tot = 0;
	getDeep(x, 0);
	sort(d+1, d+1+tot);
	int i=1, j=tot, sum=0;
	while(i<j){
		if(d[i]+d[j]<=k)	sum += j - i, i++;
		else	j--;
	}
	return sum;
}
void work(int x){
	dep[x] = 0;
	vis[x] = true;
	ans += calc(x);
	for(int i=hea[x]; i; i=edge[i].nxt){
		int t=edge[i].too;
		if(!vis[t]){
			dep[t] = edge[i].val;
			ans -= calc(t);
			sze = siz[t];
			rot = 0;
			getRoot(t, 0);
			work(rot);
		}
	}
}
int main(){
	while(scanf("%d %d", &n, &k)!=EOF){
		if(!n || !k)	break;
		init();
		for(int i=1; i<n; i++){
			scanf("%d %d %d", &uu, &vv, &ww);
			add_edge(uu, vv, ww);
			add_edge(vv, uu, ww);
		}
		gra[0] = 0x3f3f3f3f;
		sze = n;
		getRoot(1, 0);
		work(rot);
		printf("%d
", ans);
	}
	return 0;
}