poj 1741 Tree 树上的分治

http://poj.org/problem?id=1741

算法：

寻找树的重心；

分治；

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
const int maxn = 10010;
int n , m , k;
int size[maxn] , mx[maxn] , mi , dis[maxn] , center , cnt , ans;
bool vis[maxn];
struct Edge {
	int v , w , next;
	Edge() {}
	Edge(int v,int w,int next) : v(v) , w(w) , next(next) {};	
}edge[maxn<<1];
int E , head[maxn];
void init() {
	E = 0;
	memset(head,-1,sizeof(int)*(n+1));	
}
void addedge(int u,int v,int w) {
	edge[E] = Edge(v,w,head[u]); head[u] = E++;
	edge[E] = Edge(u,w,head[v]); head[v] = E++;	
}
void dfssize(int u,int fa) {
	size[u] = 1;
	mx[u] = 0;
	for(int i=head[u];i!=-1;i=edge[i].next) {
		int v = edge[i].v;
		if(vis[v] || v == fa) continue;
		dfssize(v , u);
		size[u] += size[v];
		if(size[v] > mx[u]) mx[u] = size[v];	
	}
	return;
}
void dfscenter(int r,int u,int fa) {
	int tmp = max(mx[u] , size[r]-size[u]);
	if(tmp < mi) mi = tmp , center = u;
	for(int i=head[u];i!=-1;i=edge[i].next) {
		int v = edge[i].v;
		if(vis[v] || v == fa) continue;
		dfscenter(r,v,u);	
	}	
	return;
}
void dfsdis(int u,int d,int fa) {
	dis[cnt++] = d;
	for(int i=head[u];i!=-1;i=edge[i].next) {
		int v = edge[i].v;
		if(vis[v] || v == fa) continue;
		dfsdis(v,d+edge[i].w,u);
	}
	return;
}
int calc(int u,int d) {
	int ret = 0;
	cnt = 0;
	dfsdis(u,d,-1);
	sort(dis,dis+cnt);
	int i = 0 , j = cnt - 1;
	while(i < j) {
		while(i<j & dis[i]+dis[j]>k) j--;
		ret += j-i;
		i++;
	}
	return ret;
}
void dfs(int u) {
	mi = n;
	dfssize(u,-1);
	dfscenter(u,u,-1);
	ans += calc(center,0);
	vis[center] = true;
	for(int i=head[center];i!=-1;i=edge[i].next) {
		int v = edge[i].v;
		if(vis[v]) continue;
		ans -= calc(v,edge[i].w);
		dfs(v);
	}
}
int main() {
	while(~scanf("%d%d" , &n,&k) && n+k) {
		init();
		memset(vis,false,sizeof(bool)*(n+1));
		ans = 0;
		for(int i=1;i<n;i++) {
			int u ,v , w;
			scanf("%d%d%d" ,&u,&v,&w);
			addedge(u,v,w);	
		}
		dfs(1);
		printf("%d\n" , ans);
	}
	return 0;	
}