Poj 1741 Tree (树的分治)

  1 #include <cstdio>
  2 #include <cstring>
  3 #include <iostream>
  4 #include <algorithm>
  5 using namespace std;
  6 
  7 const int maxn = 10010;
  8 struct node
  9 {
 10     int to, w, next;
 11 }edge[maxn*2];
 12 int head[maxn], vis[maxn], size[maxn], mu[maxn];
 13 int tot, n, k, ans, root, mini, d, dis[maxn];
 14 
 15 void init ()
 16 {
 17     tot = ans = 0;
 18     memset (head, -1, sizeof(head));
 19     memset (vis, 0, sizeof(vis));
 20 }
 21 void Add (int from, int to, int val)
 22 {
 23     edge[tot].to = to;
 24     edge[tot].w = val;
 25     edge[tot].next = head[from];
 26     head[from] = tot ++;
 27 }
 28 void dfsroot (int u, int father)
 29 {
 30     size[u] = 1;
 31     mu[u] = 0;
 32     for (int i=head[u]; i!=-1; i=edge[i].next)
 33     {
 34         int v = edge[i].to;
 35         if (!vis[v] && v!=father)
 36         {
 37             dfsroot (v, u);
 38             size[u] += size[v];
 39             mu[u] = max(mu[u], size[v]);
 40         }
 41     }
 42     mu[u] = max (mu[u], n-size[u]);
 43     if (mini > mu[u])
 44         root = u;
 45 }
 46 void dfsdist (int u, int father, int w)
 47 {
 48     dis[d ++] = w;
 49     for (int i=head[u]; i!=-1; i=edge[i].next)
 50     {
 51         int v = edge[i].to;
 52         if (!vis[v] && v!=father)
 53             dfsdist (v, u, w+edge[i].w);
 54     }
 55 }
 56 int calc (int u, int w)
 57 {
 58     int res = 0;
 59     d = 0;
 60     dfsdist (u, 0, w);
 61     sort (dis, dis+d);
 62     int i=0, j=d-1;
 63     while (i < j)
 64     {
 65         while (dis[i] + dis[j] > k && i<j)
 66             j --;
 67         res += j - i;
 68         i ++;
 69     }
 70     return res;
 71 }
 72 void dfs (int u)
 73 {
 74     mini = n;
 75     dfsroot (u, 0);
 76     ans += calc (root, 0);
 77     vis[root] = 1;
 78     for (int i=head[root]; i!=-1; i=edge[i].next)
 79     {
 80         int v = edge[i].to;
 81         if (!vis[v])
 82         {
 83             ans -= calc (v, edge[i].w);
 84             dfs (v);
 85         }
 86     }
 87 }
 88 int main ()
 89 {
 90     while (scanf ("%d %d", &n, &k), n+k)
 91     {
 92         init ();
 93         for (int i=1; i<n; i++)
 94         {
 95             int u, v, w;
 96             scanf ("%d %d %d", &u, &v, &w);
 97             Add (u, v, w);
 98             Add (v, u, w);
 99         }
100         dfs (1);
101         printf ("%d\n", ans);
102     }
103     return 0;
104 }