poj 1741 Tree 树上的分治

#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;
}