POJ 1987 Distance Statistics 树分治

#include <iostream>
#include <cstdio>
#include <cmath>
#include <cstring>
#include <vector>
#include <algorithm>
using namespace std;
const int N = 4e4+20, M = 1e2+10, mod = 1e9+7, inf = 1e9+1000;
typedef long long ll;

int ans, n,m,root , t = 1,K,siz[N],head[N],f[N],deep[N],d[N],allnode,vis[N];
struct edg{int to,next,v,w;}e[N * 4];
void add(int u,int v,int w) {e[t].to=v;e[t].v=w;e[t].next=head[u];head[u]=t++;}

void getroot(int x,int fa) {
    siz[x] = 1;
    f[x] = 0;
    for(int i=head[x];i;i=e[i].next) {
        int to = e[i].to;
        if(to == fa || vis[to]) continue;
        getroot(to,x);
        siz[x] += siz[to];
        f[x] = max(f[x] , siz[to]);
    }
    f[x] = max(f[x] , allnode - siz[x]);
    if(f[x] < f[root]) root = x;
}
void getdeep(int x,int fa) {
    if(d[x] <= K) deep[++deep[0]]=d[x];
    for(int i=head[x];i;i=e[i].next) {
        int to = e[i].to;
        if(to == fa || vis[to]) continue;
        d[to] = d[x] + e[i].v;
        getdeep(to,x);
    }
}
int cal(int x,int now) {
    d[x]=now;deep[0] = 0;
    getdeep(x,0);
    sort(deep+1,deep+deep[0]+1);
    int all = 0;
    for(int l=1,r=deep[0];l<r;) {
        if(deep[l]+deep[r] <= K) {all+=r-l;l++;}
        else r--;
    }
    return all;
}
void work(int x) {
    ans+=cal(x,0);
    vis[x] = 1;
    for(int i=head[x];i;i=e[i].next) {
        int to = e[i].to;
        if(vis[to]) continue;
        ans-=cal(to,e[i].v);
        allnode = siz[to];
        root = 0;
        getroot(to,root);
        work(root);
    }
}
void init()
{
    memset(head,0,sizeof(head));
    t = 1;
    ans = root = 0;
    memset(vis,0,sizeof(vis));
}
int main()
{
    while(~scanf("%d%d",&n,&m)) {
        init();
        for(int i=1;i<n;i++) {
            int a,b,c;char ch[2];
            scanf("%d%d%d%s",&a,&b,&c,ch);
            add(a,b,c) , add(b,a,c);
        }
        scanf("%d",&K);
        allnode=n;f[0]=inf;
        getroot(1,0);
        work(root);
        printf("%d\n",ans);
    }

}