poj 1741 Tree

题目大意：求树上距离<=k的点对个数

 

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define maxn 10010
#define INF 1000000000
using namespace std;
int n,k,root,sum,ans,head[maxn],num,f[maxn],vis[maxn];
int son[maxn],d[maxn],dep[maxn];
struct node{int to,pre,v;}e[maxn*2];
void Insert(int from,int to,int v){
    e[++num].to=to;
    e[num].v=v;
    e[num].pre=head[from];
    head[from]=num;
}
void getroot(int x,int father){
    son[x]=1;f[x]=0;
    for(int i=head[x];i;i=e[i].pre){
        int to=e[i].to;
        if(to==father||vis[to])continue;
        getroot(to,x);
        son[x]+=son[to];
        f[x]=max(f[x],son[to]);
    }
    f[x]=max(f[x],sum-son[x]);
    if(f[x]<f[root])root=x;
}
void getdep(int x,int father){
    dep[++dep[0]]=d[x];
    for(int i=head[x];i;i=e[i].pre){
        int to=e[i].to;
        if(to==father||vis[to])continue;
        d[to]=d[x]+e[i].v;
        getdep(to,x);
    }
}
int cal(int x,int v){
    d[x]=v;dep[0]=0;
    getdep(x,0);
    sort(dep+1,dep+dep[0]+1);
    int l=1,r=dep[0],sum=0;
    while(l<r){
        if(dep[l]+dep[r]<=k){sum+=r-l;l++;}
        else r--;
    }
    return sum;
}
void solve(int x){
    ans+=cal(x,0);
    vis[x]=1;
    for(int i=head[x];i;i=e[i].pre){
        int to=e[i].to;
        if(vis[to])continue;
        ans-=cal(to,e[i].v);
        sum=son[to];
        root=0;
        getroot(to,0);
        solve(root);
    }
}
int main(){
    while(1){
        ans=0;root=0;num=0;
        memset(vis,0,sizeof(vis));
        memset(head,0,sizeof(head));
        scanf("%d%d",&n,&k);
        if(n==0&&k==0)break;
        int x,y,z;
        for(int i=1;i<=n-1;i++){
            scanf("%d%d%d",&x,&y,&z);
            Insert(x,y,z);Insert(y,x,z);
        }
        f[0]=INF;sum=n;
        getroot(1,0);
        solve(root);
        printf("%d\n",ans);
    }
    return 0;
}