POJ1741--Tree （树的点分治） 求树上距离小于等于k的点对数

  1 #include <set>
  2 #include <map>
  3 #include <cmath>
  4 #include <ctime>
  5 #include <queue>
  6 #include <stack>
  7 #include <cstdio>
  8 #include <string>
  9 #include <vector>
 10 #include <cstdlib>
 11 #include <cstring>
 12 #include <iostream>
 13 #include <algorithm>
 14 using namespace std;
 15 typedef unsigned long long ull;
 16 typedef long long ll;
 17 const int inf = 0x3f3f3f3f;
 18 const double eps = 1e-8;
 19 const int maxn = 1e4+10;
 20 struct Edge
 21 {
 22     int to, len;
 23     Edge (int _x, int _len)
 24     {
 25         to = _x;
 26         len = _len;
 27     }
 28 };
 29 vector<Edge>G[maxn << 1];
 30 int siz[maxn];
 31 bool vis[maxn];
 32 void Get_size(int root, int father)
 33 {
 34     siz[root] = 1;
 35     for (int i = 0; i < G[root].size(); i++)
 36     {
 37         int v = G[root][i].to;
 38         if (v == father || vis[v] == true)
 39             continue;
 40         Get_size(v, root);
 41         siz[root] += siz[v];
 42     }
 43 }
 44 typedef pair<int,int>pii;
 45 pii FindGravity(int root, int father, int t)
 46 {
 47     pii res = make_pair(inf, -1);
 48     int m = 0, sum = 1;
 49     for (int i = 0; i < G[root].size(); i++)
 50     {
 51         int v = G[root][i].to;
 52         if (v == father || vis[v] == true)
 53             continue;
 54         res = min(res, FindGravity(v, root, t));
 55         m = max(m, siz[v]);
 56         sum += siz[v];
 57     }
 58     m = max(m, t-sum);
 59     res = min(res, make_pair(m, root));
 60     return res;
 61 }
 62 void Get_len(int root, int father, int d, vector<int>&len)
 63 {
 64     len.push_back(d);
 65     for (int i = 0; i < G[root].size(); i++)
 66     {
 67         int v = G[root][i].to;
 68         if (v == father || vis[v] == true)
 69             continue;
 70         Get_len(v, root, d+G[root][i].len, len);
 71     }
 72 }
 73 int K;
 74 int cnt_pair(vector<int>&ds)
 75 {
 76     int res = 0;
 77     sort (ds.begin(), ds.end());
 78     int j = ds.size() - 1;
 79     for (int i = 0; i < ds.size(); i++)
 80     {
 81         while (j > i && ds[i] + ds[j] > K)
 82         {
 83             j--;
 84         }
 85         res += (j > i ? j - i : 0);
 86     }
 87     return res;
 88 }
 89 int solve(int root)
 90 {
 91     int ans = 0;
 92     Get_size(root, -1);
 93     int s = FindGravity(root, -1, siz[root]).second;
 94     if (s == -1)
 95         return 0;
 96     vis[s] = true;
 97     for (int i = 0; i < G[s].size(); i++)
 98     {
 99         int v = G[s][i].to;
100         if (vis[v] == true)
101             continue;
102         ans += solve(v);
103     }
104     vector<int>ds;
105     ds.push_back(0);
106     for (int i = 0; i < G[s].size(); i++)
107     {
108         int v = G[s][i].to;
109         if (vis[v] == true)
110             continue;
111         vector<int>rds;
112         Get_len(v, s, G[s][i].len, rds);
113         ans -= cnt_pair(rds);
114         ds.insert(ds.end(), rds.begin(), rds.end());
115     }
116     ans += cnt_pair(ds);
117     vis[s] = false;
118     return ans;
119 }
120 int main()
121 {
122     #ifndef ONLINE_JUDGE
123         freopen("in.txt","r",stdin);
124     #endif
125     int n;
126     while ( scanf ("%d%d", &n, &K), n && K)
127     {
128         int u, v, c;
129         for (int i = 0; i <= n; i++)
130             G[i].clear();
131         for (int i = 0; i < n-1; i++)
132         {
133             scanf ("%d%d%d", &u, &v, &c);
134             G[u].push_back(Edge(v,c));
135             G[v].push_back(Edge(u,c));
136         }
137         printf("%d\n", solve(n/2+1));
138     }
139     return 0;
140 }