POJ 1741 男人八题——树分治

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>


using namespace std;

#define N 10005

struct Node {
    int v,l;
};

vector<Node> g[N];

int n,k;

int s[N],f[N];
bool done[N];
int size;
int root;
void getroot(int u,int fa) {
    s[u] = 1;
    f[u] = 0;

    for(int i=0; i < (int)g[u].size(); i++) {
        int v = g[u][i].v;
        if(v!=fa&&!done[v]) {
            getroot(v,u);
            s[u] +=s[v];
            f[u] = max(f[u],s[v]);
        }
    }

    f[u] = max(f[u],size-s[u]);
    if(f[u]<f[root]) root = u;

}

vector<int> dep;
int d[N];
void getdep(int u,int fa) {

    dep.push_back(d[u]);
    s[u] = 1;

    for(int i=0; i < (int)g[u].size(); i++) {
        int v = g[u][i].v;
        if(v!=fa&&!done[v]) {
            d[v] = d[u] + g[u][i].l;
            getdep(v,u);
            s[u]+=s[v];
        }

    }

}

int calc(int u,int init) {
    dep.clear();
    d[u] = init;
    getdep(u,0);

    sort(dep.begin(),dep.end());

    int ret = 0;
    for(int l=0,r=dep.size()-1;l<r;)
        if(dep[l]+dep[r]<=k) ret += r-l++;
        else r--;
    return ret;
}

int ans;
void work(int u) {

    ans+=calc(u,0);
    done[u] = true;

    for(int i=0; i < (int)g[u].size(); i++) {
        int v = g[u][i].v;
        if(!done[v]) {
            ans-=calc(v,g[u][i].l);
            f[0] = size = s[v];
            getroot(v,root=0);
            work(root);
        }
    }

}

int main()
{
    while(scanf("%d%d",&n,&k),n) {

        for(int i=0; i <= n; i++) g[i].clear();

        memset(done,0,sizeof(done));

        int u,v,l;
        for(int i=1; i < n; i++) {
            scanf("%d%d%d",&u,&v,&l);
            g[u].push_back((Node){v,l});
            g[v].push_back((Node){u,l});
        }

        f[0] = size = n;
        getroot(1,root=0);
        ans = 0;
        work(root);

        printf("%d\n",ans);

    }
    return 0;
}