POJ 1741 树上的点分治

题目大意：

找到树上点对间距离不大于K的点对数

这是一道简单的练习点分治的题，注意的是为了防止点分治时出现最后分治出来一颗子树为一条直线，所以用递归的方法求出最合适的root点

 #include <cstdio>
 #include <cstring>
 #include <iostream>
 #include <vector>
 #include <algorithm>

 using namespace std;
 #define N 10005
 int n , m , k , first[N];

 struct Edge{
     int y , next , d;
     Edge(){}
     Edge(int y , int next , int d):y(y),next(next),d(d){}
 }e[N<<];

 void add_edge(int x , int y , int d)
 {
     e[k] = Edge(y , first[x] , d);
     first[x] = k++;
 }

 int sz[N] , dis[N] , f[N] , d[N] , cnt , root , ret;
 bool use[N];

 void find_root(int u , int fa , int size)
 {
     sz[u] =  , f[u] = ;
     int v;
     for(int i=first[u] ; ~i ; i=e[i].next){
         if(use[v=e[i].y] || v==fa) continue;
         find_root(v , u , size);
         sz[u] += sz[v] ;
         f[u] = max(f[u] , sz[v]);
     }
     f[u] = max(f[u] , size-sz[u]);
     if(f[u]<f[root]) root = u;
 }

 void dfs(int u , int fa)
 {
     d[cnt++] = dis[u];
     sz[u] = ;
     int v;
     for(int i=first[u] ; ~i ; i=e[i].next){
         if(use[v=e[i].y] || v==fa) continue;
         dis[v] = dis[u]+e[i].d;
         if(dis[v]>m) continue;
         dfs(v , u);
         sz[u] += sz[v];
     }
 }

 int cal(int u , int val)
 {
     dis[u] = val , cnt=;
     dfs(u , );
     sort(d , d+cnt);
     int ret = ;
     for(int l= , r=cnt- ; l<r ; )
         if(d[l]+d[r]<=m) ret+=r-l++;
         else r--;
     return ret;
 }

 void solve(int u)
 {
     ret+=cal(u , );
     use[u] =true;
     int v;
     for(int i=first[u] ; ~i ; i=e[i].next){
         if(use[v=e[i].y]) continue;
         ret -= cal(v , e[i].d);
         find_root(v , root= , sz[v]);
         solve(root);
     }
 }

 int main()
 {
    // freopen("in.txt" , "r" , stdin);
     int x,y,d;
     while(scanf("%d%d" , &n , &m) , n+m)
     {
         memset(first , - , sizeof(first));
         k = ;
         for(int i= ; i<n ; i++){
             scanf("%d%d%d" , &x , &y , &d);
             add_edge(x , y , d);
             add_edge(y , x , d);
         }
         memset(use ,  , sizeof(use));
         ret= , f[] = 1e9;
         find_root( , root= , n);
         solve(root);
         printf("%d\n" , ret);
     }
 }