POJ 1741 Tree

 1 #include <iostream>
 2 #include <cstdio>
 3 #include <algorithm>
 4 #include <cstring>
 5 using namespace std;
 6 const int maxn = 100010;
 7 struct arc{
 8     int to,w,next;
 9     arc(int x = 0,int y = 0,int z = -1){
10         to = x;
11         w = y;
12         next = z;
13     }
14 }e[maxn<<1];
15 bool vis[maxn];
16 int head[maxn],d[maxn],sz[maxn],maxson[maxn],cnt,tot;
17 int n,k;
18 void add(int u,int v,int w){
19     e[tot] = arc(v,w,head[u]);
20     head[u] = tot++;
21 }
22 void dfs(int u,int fa){
23     sz[u] = 1;
24     maxson[u] = 0;
25     for(int i = head[u]; ~i; i = e[i].next){
26         if(e[i].to == fa || vis[e[i].to]) continue;
27         dfs(e[i].to,u);
28         sz[u] += sz[e[i].to];
29         maxson[u] = max(maxson[u],sz[e[i].to]);
30     }
31 }
32 int FindRoot(int sum,int u,int fa){
33     int ret = u;
34     for(int i = head[u]; ~i; i = e[i].next){
35         if(e[i].to == fa || vis[e[i].to]) continue;
36         int x = max(sum - sz[ret],maxson[ret]);
37         int y = FindRoot(sum,e[i].to,u);
38         int z = max(sum - sz[y],maxson[y]);
39         if(z < x) ret = y;
40     }
41     return ret;
42 }
43 void dis(int u,int fa,int depth){
44     d[cnt++] = depth;
45     for(int i = head[u]; ~i; i = e[i].next){
46         if(e[i].to == fa || vis[e[i].to] || depth + e[i].w > k) continue;
47         dis(e[i].to,u,depth + e[i].w);
48     }
49 }
50 int calc(int u,int fa,int depth){
51     int ret = cnt = 0;
52     dis(u,fa,depth);
53     sort(d, d + cnt);
54     int low = 0,high = cnt - 1;
55     while(low < high){
56         if(d[low]  + d[high] <= k) ret += high - low++;
57         else --high;
58     }
59     return ret;
60 }
61 int solve(int u,int fa){
62     dfs(u,0);
63     int root = FindRoot(sz[u],u,0);
64     vis[root] = true;
65     int ret = calc(root,0,0);
66     for(int i = head[root]; ~i; i = e[i].next){
67         if(vis[e[i].to]) continue;
68         ret -= calc(e[i].to,root,e[i].w);
69         ret += solve(e[i].to,root);
70     }
71     return ret;
72 }
73 int main(){
74     int u,v,w;
75     while(scanf("%d%d",&n,&k),n||k){
76         memset(head,-1,sizeof head);
77         memset(vis,false,sizeof vis);
78         tot = 0;
79         for(int i = 1; i < n; ++i){
80             scanf("%d%d%d",&u,&v,&w);
81             add(u,v,w);
82             add(v,u,w);
83         }
84         printf("%d\n",solve(1,0));
85     }
86     return 0;
87 }