树点分治入门题poj1741

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int N=11111;
const int M=55555;
const int INF=1e9;
struct node{
    int v,next,w;
}e[M];
int head[N],tot;
int n,k,vis[N],ans,root,num;
void init(){
    memset(vis,0,sizeof(vis));
    memset(head,-1,sizeof(head));
    tot=ans=0;
}
void add(int u,int v,int w){
   e[tot].v=v;e[tot].w=w;e[tot].next=head[u];head[u]=tot++;
}
int mx[N],size[N],mi,dis[N];
void dfssize(int u,int fa){
    size[u]=1;
    mx[u]=0;
    for(int i=head[u];~i;i=e[i].next){
        int v=e[i].v;
        if(v!=fa&&!vis[v]) {
            dfssize(v,u);
            size[u]+=size[v];
            if(size[v]>mx[u]) mx[u]=size[v];
        }
    }
}
void dfsroot(int r,int u,int fa){
   if(size[r]-size[u]>mx[u]) mx[u]=size[r]-size[u];
   if(mx[u]<mi) mi=mx[u],root=u;
   for(int i=head[u];~i;i=e[i].next){
    int v=e[i].v;
    if(v!=fa&&!vis[v]) dfsroot(r,v,u);
   }
}
void dfsdis(int u,int d,int fa){
    dis[num++]=d;
    for(int i=head[u];~i;i=e[i].next){
        int v=e[i].v;
        if(v!=fa&&!vis[v]) dfsdis(v,d+e[i].w,u);
    }
}
int calc(int u,int d){
   int ret=0;
   num=0;
   dfsdis(u,d,0);
   sort(dis,dis+num);
   int i=0,j=num-1;
   while(i<j){
    while(dis[i]+dis[j]>k&&i<j) --j;
    ret+=j-i;
    ++i;
   }
   return ret;
}
void dfs(int u){//由于每次都取重心，所以最糟糕的情况是logN层,然后每层小于等于N个数遍历，所以复杂度N*logN
    mi=n;
    dfssize(u,0);
    dfsroot(u,u,0);
    ans+=calc(root,0);
    vis[root]=1;
    for(int i=head[root];~i;i=e[i].next){
        int v=e[i].v;
        if(!vis[v]){
         ans-=calc(v,e[i].w);
        dfs(v);
        }
    }
}
int main(){
    while(scanf("%d%d",&n,&k)!=EOF){
        if(!n&&!k) break;
        init();
        int u,v,w;
        for(int i=0;i<n-1;++i) {
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        dfs(1);
        printf("%d\n",ans);
    }
}