POJ 1741 树上的点分治

 1 #include <cstdio>
 2 #include <cstring>
 3 #include <iostream>
 4 #include <vector>
 5 #include <algorithm>
 6 
 7 using namespace std;
 8 #define N 10005
 9 int n , m , k , first[N];
10 
11 struct Edge{
12     int y , next , d;
13     Edge(){}
14     Edge(int y , int next , int d):y(y),next(next),d(d){}
15 }e[N<<1];
16 
17 void add_edge(int x , int y , int d)
18 {
19     e[k] = Edge(y , first[x] , d);
20     first[x] = k++;
21 }
22 
23 int sz[N] , dis[N] , f[N] , d[N] , cnt , root , ret;
24 bool use[N];
25 
26 void find_root(int u , int fa , int size)
27 {
28     sz[u] = 1 , f[u] = 0;
29     int v;
30     for(int i=first[u] ; ~i ; i=e[i].next){
31         if(use[v=e[i].y] || v==fa) continue;
32         find_root(v , u , size);
33         sz[u] += sz[v] ;
34         f[u] = max(f[u] , sz[v]);
35     }
36     f[u] = max(f[u] , size-sz[u]);
37     if(f[u]<f[root]) root = u;
38 }
39 
40 void dfs(int u , int fa)
41 {
42     d[cnt++] = dis[u];
43     sz[u] = 1;
44     int v;
45     for(int i=first[u] ; ~i ; i=e[i].next){
46         if(use[v=e[i].y] || v==fa) continue;
47         dis[v] = dis[u]+e[i].d;
48         if(dis[v]>m) continue;
49         dfs(v , u);
50         sz[u] += sz[v];
51     }
52 }
53 
54 int cal(int u , int val)
55 {
56     dis[u] = val , cnt=0;
57     dfs(u , 0);
58     sort(d , d+cnt);
59     int ret = 0;
60     for(int l=0 , r=cnt-1 ; l<r ; )
61         if(d[l]+d[r]<=m) ret+=r-l++;
62         else r--;
63     return ret;
64 }
65 
66 void solve(int u)
67 {
68     ret+=cal(u , 0);
69     use[u] =true;
70     int v;
71     for(int i=first[u] ; ~i ; i=e[i].next){
72         if(use[v=e[i].y]) continue;
73         ret -= cal(v , e[i].d);
74         find_root(v , root=0 , sz[v]);
75         solve(root);
76     }
77 }
78 
79 int main()
80 {
81    // freopen("in.txt" , "r" , stdin);
82     int x,y,d;
83     while(scanf("%d%d" , &n , &m) , n+m)
84     {
85         memset(first , -1 , sizeof(first));
86         k = 0;
87         for(int i=1 ; i<n ; i++){
88             scanf("%d%d%d" , &x , &y , &d);
89             add_edge(x , y , d);
90             add_edge(y , x , d);
91         }
92         memset(use , 0 , sizeof(use));
93         ret=0 , f[0] = 1e9;
94         find_root(1 , root=0 , n);
95         solve(root);
96         printf("%d\n" , ret);
97     }
98 }