poj 1741 树的点分治

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 12913		Accepted: 4128

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

Source

LouTiancheng@POJ

#include <iostream>
#include<stdio.h>
#include<string.h>
#include<algorithm>
#include<math.h>
using namespace std;
int n,k;
int ans;
const int maxn=11000;
const int maxm=21111;
struct node
{
    int v;
    int w;
    int next;
}e[maxm];
int pre[maxn],cnt;
int data[maxn];
int dis[maxn];
int vis[maxn];
struct centertree
{
    int n;
    int ans;
    int son[maxn];
    int size;
    void dfs(int u,int fa)
    {
        son[u]=1;
        int res=-1;
        for(int i=pre[u];i!=-1;i=e[i].next)
        {
            int v=e[i].v;
            if(v==fa) continue;
            if(vis[v]) continue;
            dfs(v,u);
            son[u]+=son[v];
            res=max(res,son[v]);
        }
        res=max(res,n-son[u]);
         if(res<size)
         {
             ans=u;
             size=res;
         }
    }
    int getcenter(int x)
    {
        ans=0;
        size=1e9;
        dfs(x,-1);
        return ans;
    }


}cent;
void add(int u,int v,int w)
{
    e[cnt].v=v;
    e[cnt].w=w;
    e[cnt].next=pre[u];
    pre[u]=cnt++;
}
int len;
void getarray(int u,int fa)
{
    data[++len]=dis[u];
    for(int i=pre[u];i!=-1;i=e[i].next)
    {
        int v=e[i].v;
        int w=e[i].w;
        if(v==fa) continue;
        if(vis[v]) continue;
        dis[v]=dis[u]+w;
       getarray(v,u);
    }
}
int calc(int u,int now)
{
    dis[u]=now;
    len=0;
    getarray(u,-1);
    sort(data+1,data+1+len);
    int res=0;
    int l=1,r=len;
    while(l<r)
    {
        if(data[l]+data[r]<=k)
        {
            res+=(r-l);
            l++;
        }
        else
            r--;
    }
    return res;
}
void solve(int u)
{
    ans+=calc(u,0);
    vis[u]=1;
    for(int i=pre[u];i!=-1;i=e[i].next)
    {
        int v=e[i].v;
        int w=e[i].w;
        if(vis[v]) continue;
        ans-=calc(v,w);
        cent.n=cent.son[v];
        int rt=cent.getcenter(v);
        solve(rt);
    }
}
int main()
{
    while(~scanf("%d%d",&n,&k))
    {
        if(n==0&&k==0)
            break;
            cnt=0;
            memset(pre,-1,sizeof(pre));
            memset(vis,0,sizeof(vis));
        for(int i=1;i<n;i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        ans=0;
        cent.n=n;
        int root=cent.getcenter(1);
        solve(root);
        printf("%d\n",ans);
    }

}