poj1741 Tree点分治

点分治入门题.
每次统计完到重心的dis之后sort一遍统计答案.

#include<stdio.h>
#include<cstring>
#include<algorithm>
using namespace std;
const int maxn=10005;
const int inf=210000000;
bool vis[maxn];
int h[3*maxn],f[maxn],n,k,deep[maxn],dis[maxn],ans,root,num,siz[maxn],sum;
inline const int read(){ 
    register int f=1,x=0;
    register char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=(x<<3)+(x<<1)+ch-'0';ch=getchar();}
    return x*f;
}
struct edge{
    int v,nxt,w;
}e[3*maxn];
void add(int u,int v,int w){
    e[++num].v=v;
    e[num].w=w;
    e[num].nxt=h[u];
    h[u]=num;
}
void get_heavy(int u,int fa){
    siz[u]=1,f[u]=0;
    for(int i=h[u];i;i=e[i].nxt){
        int v=e[i].v;
        if(v==fa||vis[v]) continue; 
        get_heavy(v,u);
        siz[u]+=siz[v];
        f[u]=max(f[u],siz[v]);
    }
    f[u]=max(f[u],sum-siz[u]);
    if(f[u]<f[root]) root=u;
}
void get_deep(int u,int fa){
    deep[++deep[0]]=dis[u];
    for(int i=h[u];i;i=e[i].nxt){
        int v=e[i].v;
        if(v==fa||vis[v]) continue;
        dis[v]=dis[u]+e[i].w;
        get_deep(v,u);
    }
}
inline int cal(int u,int now){
    dis[u]=now,deep[0]=0;
    get_deep(u,0);
    sort(deep+1,deep+deep[0]+1);
    int t=0,l,r;
    for(l=1,r=deep[0];l<r;){
     if(deep[l]+deep[r]<=k) t+=r-l,l++;
     else r--;
    }
    return t;
}
void solve(int u){
    ans+=cal(u,0);
    vis[u]=true;
    for(int i=h[u];i;i=e[i].nxt){
        int v=e[i].v;
        if(vis[v]) continue;
        ans-=cal(v,e[i].w);
        sum=siz[v];
        root=0;
        get_heavy(v,root);
        solve(root);
    }
}
int main(){
    while(1){
       ans=0,root=0,num=0;
       memset(vis,false,sizeof(vis));
       memset(h,0,sizeof(h));
       n=read(),k=read();
       if(n==0&&k==0) break;
       int x,y,w;
       for(int i=1;i<+n;i++){
           x=read(),y=read(),w=read();
           add(x,y,w),add(y,x,w);
       }
       f[0]=inf,sum=n;
       get_heavy(1,0);
       solve(root);
       printf("%d\n",ans);
    }
}