POJ 1741 Tree （树的点分治入门）

view code#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <queue>
using namespace std;
const int N = 10010;
int n, k, pre[N], ans, mi, rt, siz[N], num;
bool vis[N];

struct edge
{
    int u, v, w, next;
    edge() {}
    edge(int u, int v, int w, int next):u(u),v(v),w(w),next(next) {}
}e[N<<1];
int ecnt;

void init()
{
    ans = ecnt = 0;
    memset(pre, -1, sizeof(pre));
    memset(vis, 0, sizeof(vis));
}

inline void add(int u, int v, int w)
{
    e[ecnt] = edge(u, v, w, pre[u]);
    pre[u] = ecnt++;
}

void getroot(int u, int fa)
{
    siz[u] = 1;
    int mx = 0;
    for(int i=pre[u]; ~i; i=e[i].next)
    {
        int v = e[i].v;
        if(v==fa || vis[v]) continue;
        getroot(v, u);
        siz[u] += siz[v];
        mx = max(mx, siz[v]);
    }
    mx = max(mx, siz[0]-siz[u]);
    if(mx <mi) mi = mx, rt = u;
}

int dis[N];
void getdis(int u, int d, int fa)
{
    dis[num++] = d;
    for(int i=pre[u]; ~i; i=e[i].next)
    {
        int v = e[i].v;
        if(v==fa || vis[v]) continue;
        getdis(v, d+e[i].w, u);
    }
}

int calc(int u, int d)
{
    int res = 0;
    num = 0;
    getdis(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--;
        res += j-i;
        i++;
    }
    return res;
}

void solve(int u, int cnt)
{
    mi = n;
    siz[0] = cnt;
    getroot(u, 0);
    ans += calc(rt, 0);
    vis[rt] = 1;
    for(int i=pre[rt]; ~i; i=e[i].next)
    {
        int v = e[i].v;
        if(vis[v]) continue;
        ans -= calc(v, e[i].w);
        solve(v, siz[v]);
    }

}

int main()
{
//    freopen("in.txt", "r", stdin);
    while(scanf("%d%d", &n, &k)>0 && (n|k))
    {
        int u, v, w;
        init();
        for(int i=1; i<n; i++)
        {
            scanf("%d%d%d", &u, &v, &w);
            add(u, v, w);
            add(v, u, w);
        }
        solve(1, n);
        printf("%d\n", ans);
    }
    return 0;
}