POJ1741Tree [点分治]【学习笔记】

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N=10005,INF=1e9+5;
inline int read(){
    char c=getchar();int x=0,f=1;
    while(c<‘0‘||c>‘9‘){if(c==‘-‘)f=-1;c=getchar();}
    while(c>=‘0‘&&c<=‘9‘){x=x*10+c-‘0‘;c=getchar();}
    return x*f;
}
int n,L,u,v,w;
struct edge{
    int v,w,ne;
}e[N<<1];
int h[N],cnt;
inline void ins(int u,int v,int w){
    cnt++;
    e[cnt].v=v;e[cnt].w=w;e[cnt].ne=h[u];h[u]=cnt;
    cnt++;
    e[cnt].v=u;e[cnt].w=w;e[cnt].ne=h[v];h[v]=cnt;
}

int size[N],d[N],vis[N],root,sum;
void dfsRoot(int u,int fa){
    size[u]=1;d[u]=0;
    for(int i=h[u];i;i=e[i].ne){
        int v=e[i].v;
        if(vis[v]||v==fa) continue;
        dfsRoot(v,u);
        size[u]+=size[v];
        d[u]=max(d[u],size[v]);
    }
    d[u]=max(d[u],sum-size[u]);
    if(d[u]<d[root]) root=u;
}
int deep[N],a[N];
void dfsDeep(int u,int fa){
    a[++a[0]]=deep[u];
    for(int i=h[u];i;i=e[i].ne){
        int v=e[i].v;
        if(vis[v]||v==fa) continue;
        deep[v]=deep[u]+e[i].w;
        dfsDeep(v,u);
    }
}

int cal(int u,int now){
    deep[u]=now;a[0]=0;
    dfsDeep(u,0);
    sort(a+1,a+1+a[0]);
    int l=1,r=a[0],ans=0;
    while(l<r){
        if(a[l]+a[r]<=L) ans+=r-l,l++;
        else r--;
    }
    return ans;
}
int ans;
void dfsSol(int u){//printf("dfs %d\n",u);
    vis[u]=1;
    ans+=cal(u,0);
    for(int i=h[u];i;i=e[i].ne){
        int v=e[i].v;
        if(vis[v]) continue;
        ans-=cal(v,e[i].w);
        sum=size[v];
        root=0;dfsRoot(v,0);
        dfsSol(root);
    }
}

int main(){
    //freopen("in.txt","r",stdin);
    while(true){
        n=read();L=read();if(n==0) break;
        cnt=0;memset(h,0,sizeof(h));
        memset(vis,0,sizeof(vis));
        ans=0;
        for(int i=1;i<=n-1;i++) u=read(),v=read(),w=read(),ins(u,v,w);
        sum=n;
        root=0;d[0]=INF;
        dfsRoot(1,0);
        dfsSol(root);
        printf("%d\n",ans);
    }
}