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);
}

}