POJ-1741 点分治

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 18156		Accepted: 5930

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

题目链接：

http://poj.org/problem?id=1741

题目大意：

        计算树上距离不超过k的点对数目，n<=10000。

解题思路：

        点分治模板题。

AC代码：

import java.util.*;

public class Main {
	static int n,k,a,b,c,ee,ans,ct;
	static int[] mm=new int[10005];
	static int[] mx=new int[10005];
	static int[] dis=new int[10005];
	static int[] size=new int[10005];
	static int[] head=new int[10005];
	static int[] edge=new int[20005];
	static int[] next=new int[20005];
	static int[] cost=new int[20005];
	static void add(int a,int b,int c)
	{
		edge[ee]=b;next[ee]=head[a];
		cost[ee]=c;head[a]=ee++;
	}
	static void dfs(int u)
	{
		getsize(u,0);
		int rt=getroot(u,0,u);
		ans+=cal(rt,0);
		mm[rt]=1;int v;
		for(int i=head[rt];i!=-1;i=next[i])
		{
			v=edge[i];
			if(mm[v]==1) continue;
			ans-=cal(v,cost[i]);
			dfs(v);
		}
	}
	static int cal(int u,int d)
	{
		int res=0;ct=0;
		getdis(u,0,d);
		Arrays.sort(dis,1,ct+1);
		int i=1,j=ct;
		while(i<j)
		{
			while(dis[i]+dis[j]>k&&i<j)
				j--;
			res+=j-i;i++;
		}
		return res;
	}
	static void getdis(int u,int fa,int d)
	{
		dis[++ct]=d;int v;
		for(int i=head[u];i!=-1;i=next[i])
		{
			v=edge[i];
			if(mm[v]==1||v==fa) continue;
			getdis(v,u,d+cost[i]);
		}
	}
	static void getsize(int u,int fa)
	{
		size[u]=1;mx[u]=0;
		int v;
		for(int i=head[u];i!=-1;i=next[i])
		{
			v=edge[i];
			if(mm[v]==1||v==fa) continue;
			getsize(v,u);
			size[u]+=size[v];
			mx[u]=Math.max(mx[u],size[v]);
		}
	}
	static int getroot(int u,int fa,int r)
	{
		int res=u,v,p;
		mx[u]=Math.max(mx[u],size[r]-size[u]);
		for(int i=head[u];i!=-1;i=next[i])
		{
			v=edge[i];
			if(mm[v]==1||v==fa) continue;
			p=getroot(v,u,r);
			if(mx[p]<mx[res]) res=p;
		}
		return res;
	}

	public static void main(String[] args) {
		Scanner in=new Scanner(System.in);
		while(in.hasNext())
		{
			n=in.nextInt();k=in.nextInt();
			if(n==0&&k==0) break;
			Arrays.fill(head,-1);ee=0;
			for(int i=1;i<n;i++)
			{
				a=in.nextInt();b=in.nextInt();
				if(a>b) { c=a;a=b;b=c;}
				c=in.nextInt();
				add(a,b,c);add(b,a,c);
			}
			Arrays.fill(mm,0);ans=0;
			dfs(1);
			System.out.println(ans);
		}
	}
}