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

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;

}