POJ1741 Tree （点分治）

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 25772		Accepted: 8566

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001). 
Define dist(u,v)=The min distance between node u and v. 
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k. 
Write a program that will count how many pairs which are valid for a given tree. 

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l. 
The last test case is followed by two zeros. 

Output

For each test case output the answer on a single line.

Sample Input

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

Sample Output

8

code

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

using namespace std;

const int N = 10100;

struct Edge{
    int to,nxt,w;
}e[N<<1];
int head[N],son[N],f[N],deth[N],d[N];
bool vis[N];
int ans,tot,Root,sum,n,k;

void init() {
    memset(head,0,sizeof(head));
    memset(vis,false,sizeof(vis));
    ans = tot = Root = 0;
}
void add_edge(int u,int v,int w) {
    e[++tot].to = v;e[tot].w = w;e[tot].nxt = head[u];head[u] = tot;
}
void getroot(int u,int pa) {
    son[u] = 1;f[u] = 0;
    for (int i=head[u]; i; i=e[i].nxt) {
        int v = e[i].to;
        if (v==pa||vis[v]) continue;
        getroot(v,u);
        son[u] += son[v];
        f[u] = max(f[u],son[v]);
    }
    f[u] = max(f[u],sum-son[u]);
    if (f[u] < f[Root]) Root = u;
}
void getdeth(int u,int pa) {
    deth[++deth[0]] = d[u];
    for (int i=head[u]; i; i=e[i].nxt) {
        int v = e[i].to,w = e[i].w;
        if (v==pa||vis[v]) continue;
        d[v] = d[u] + w;
        getdeth(v,u);
    }
}
int calcc(int u,int w) {
    deth[0] = 0;
    d[u] = w;
    getdeth(u,0);
    sort(deth+1,deth+deth[0]+1);
    int l = 1,r = deth[0],ret = 0;
    while (l < r) {
        if (deth[l] + deth[r] <= k) ret += r-l,l++;
        else r--;
    }
    return ret;
}
void work(int u) {
    ans += calcc(u,0);
    vis[u] = 1;
    for (int i=head[u]; i; i=e[i].nxt) {
        int v = e[i].to,w = e[i].w;
        if (vis[v]) continue;
        ans -= calcc(v,w);
        sum = son[v];
        Root = 0;
        getroot(v,0);
        work(Root);
    }
}
int main() {
    while (scanf("%d%d",&n,&k)!=EOF && !(n==0&&k==0)) {
        init();
        for (int a,b,c,i=1; i<n; ++i) {
            scanf("%d%d%d",&a,&b,&c);
            add_edge(a,b,c);add_edge(b,a,c);
        }
        f[0] = 1e9;
        sum = n;
        getroot(1,0);
        work(Root);
        printf("%d
",ans);
    }
    return 0;
}