poj1741(点分治、树上点对）

题目大意：一颗树，问由多少个点对（u，v)满足dis(u,v)<=k

        树的点分治。。。求先求以重心为根的树中满足dis(u,v)<=k的点对数，减去重心为根的所有子树中满足条件的点对数，即为路径经过重心的点对数，加起来即为答案。

      每次对以重心为根的树进行dfs，子树大小减半，再以同样的方式每棵dfs子树，所以递归的深度为logn。每次查找答案都要遍历完其子树每个点，所以对于每一层递归遍历的结点总数为剩余没访问过的点的总和，为n。由于遍历完子树所有结点后要对其深度排序，所以总时间复杂度为O(n*logn*logn)

     代码

#include<iostream>
#include<stdio.h>
#include<string.h>
#include<algorithm>
using namespace std;
int m,root,siz[20010],dep[20010],vis[20010],ma[20010],A[20010],cnt,k;
int e,p[20010],nex[20010],head[20010],w[20010];
void add(int a,int b,int c){
     p[e]=b;
     nex[e]=head[a];
     head[a]=e;
     w[e++]=c;
}
void dfsroot(int u,int fa,int sum){
    int i,v;
    ma[u]=0;
    for(i=head[u];i;i=nex[i]){
        v=p[i];
        if(v!=fa&&vis[v]==0){
            dfsroot(v,u,sum);
            if(ma[u]<siz[v]) ma[u]=siz[v];
        }
    }
    if(ma[u]<sum-siz[u]) ma[u]=sum-siz[u];
    if(ma[u]<ma[root]||root==0) root=u;
}
void dfsdep(int u,int fa){
    int i,v;
    A[k++]=dep[u];
    for(i=head[u];i;i=nex[i]){
        v=p[i];
        if(vis[v]==0&&v!=fa){
            dep[v]=dep[u]+w[i];
            dfsdep(v,u);
        }
    }
}
void dfssiz(int u,int fa){
    int i,v;
    siz[u]=1;
    for(i=head[u];i;i=nex[i]){
        v=p[i];
        if(vis[v]==0&&v!=fa){
            dfssiz(v,u);
            siz[u]+=siz[v];
        }
    }
}
int cal(int u){
    k=1;
    dfsdep(u,0);
    sort(A+1,A+k);
    int ans=0,l=1,r=k-1;
    while(l<r){
        while(l<r&&A[l]+A[r]>m) r--;
        ans+=r-l;
        l++;
    }
    return ans;
}
void dfs(int u,int fa){
    int i,v;
    dfssiz(u,0);
    root=0;
    dfsroot(u,0,siz[u]);
    vis[root]=1;
    dep[root]=0;
    cnt+=cal(root);
    for(i=head[root];i;i=nex[i]){
        v=p[i];
        if(vis[v]==0&&v!=fa){
            dep[v]=w[i];
            cnt-=cal(v);
            dfs(v,root);
        }
    }
}
int main(){
    int n,i,a,b,c;
    while(scanf("%d%d",&n,&m)!=EOF){
        if(n==0&&m==0) break;
        memset(head,0,sizeof(head));
        memset(vis,0,sizeof(vis));
        e=1;
        for(i=1;i<n;i++){
            scanf("%d%d%d",&a,&b,&c);
            add(a,b,c);
            add(b,a,c);
        }
        cnt=0;
        dfs(1,0);
        printf("%d\n",cnt);
    }
    return 0;
}