POJ 1741 tree

 1 #include <cstdio>
 2 #include <algorithm>
 3 #include <vector>
 4 #include <cstring>
 5 #include <iostream>
 6 using namespace std;
 7 #define N 10010
 8 struct node {
 9     int to, w;
10     node() {};
11     node(int _v, int _l): to(_v), w(_l) {};
12 };
13 vector<node> g[N];
14 int n, k, sizee, s[N], f[N], root, d[N], K, ans;
15 vector<int>dep;
16 bool vis[N];
17 void getroot(int now, int fa) {
18     s[now] = 1; f[now] = 0;
19     int u;
20     for (int i=0; i<g[now].size(); i++)
21         if ((u = g[now][i].to) != fa && !vis[u]) {
22             getroot(u, now);
23             s[now] += s[u];
24             f[now] = max(f[now], s[u]);
25         }
26     f[now] = max(f[now], sizee-s[now]);
27     if (f[now] < f[root]) root = now;
28 }
29 void getdep(int now, int fa) {
30     dep.push_back(d[now]);
31     int u;
32     for (int i=0; i<g[now].size(); i++)
33         if ((u = g[now][i].to) != fa && !vis[u]) {
34             d[u] = d[now] + g[now][i].w;
35             getdep(u, now);
36         }
37 }
38 int calc(int now, int init) {
39     dep.clear(); d[now] = init;
40     getdep(now, 0);
41     sort(dep.begin(), dep.end());
42     int ret = 0;
43     for (int l=0, r=dep.size()-1; l<r; )
44         if (dep[l] + dep[r] <= K) ret += r-l++;
45         else r--;
46     return ret;
47 }
48 void work(int now) {
49     int u;
50     ans += calc(now, 0);
51     vis[now] = true;
52     for (int i=0; i<g[now].size(); i++)
53         if (!vis[u = g[now][i].to]) {
54             ans -= calc(u, g[now][i].w);
55             f[0] = sizee = s[u];
56             getroot(u, root=0);
57             work(root);
58         }
59 }
60 int main() {
61     while (scanf("%d%d", &n, &K) == 2) {
62         if (n == 0 && K == 0) break;
63         for (int i=0; i<=n; i++) g[i].clear();
64         memset(vis, false, sizeof(vis));
65 
66         int u, v, l;
67         for (int i=1; i<n; i++) {
68             scanf("%d%d%d", &u, &v, &l);
69             g[u].push_back(node(v, l));
70             g[v].push_back(node(u, l));
71         }
72         f[0] = sizee = n;
73         getroot(1, root=0);
74         ans = 0;
75         work(root);
76         printf("%d\n", ans);
77     }
78     return 0;
79 }