POJ1741——Tree (树分治之点分治)

时间:2020-12-25 04:23:07
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
代码:
#include <cstdio>
#include <algorithm>
#include <cstring>
#include <cmath>

using namespace std;

const int INF = 0x3f3f3f3f;
const int MAXN = 10010;

bool vis[MAXN];
int head[MAXN];
int son[MAXN];//记录当前以v为根的树的大小。 
int f[MAXN];//f记录当前以v为根的树的最大子树的大小。 
int d[MAXN];//存每个点到root的距离。 
int deep[MAXN];//同样存是每个点到root的距离,用于sort。 
struct Edge{
	int next;
	int to;
	int value;
}E[MAXN*2];//由于每条边存两遍,所以要两倍空间。 
int N,K;
int top,ans,root,sum;
//从左到右分别为存边数组的指针,结果,当前树最优根,当前树大小。 

void add_edge(int from,int to,int value){//邻接表加入新边。 
	E[++top].next = head[from];
	head[from] = top;
	E[top].to = to;
	E[top].value = value;
}

void getroot(int v,int father){//找最优根(重心) 
	son[v] = 1;//感开始就根节点一个点。 
	f[v] = 0;
	for(int i=head[v] ; i ; i=E[i].next){
		if(E[i].to!=father && !vis[E[i].to]){
			getroot(E[i].to,v);
			son[v] += son[E[i].to];
			f[v] = max(f[v],son[E[i].to]);//与子树比较 
		}
	}
	f[v] = max(f[v],sum-son[v]);//别忘了以v父节点为根的树其实也是子树。 
	if(f[v]<f[root])root = v;//更新当前重心
}

void getdeep(int v,int father){
	deep[++deep[0]] = d[v];
	for(int i=head[v] ; i ; i=E[i].next){
		if(E[i].to!=father && !vis[E[i].to]){
			d[E[i].to] = d[v] + E[i].value;
			getdeep(E[i].to,v);
		}
	}
}

int cal(int v,int cost){
	d[v] = cost;
	deep[0] = 0;
	getdeep(v,0);
	sort(deep+1,deep+deep[0]+1);//对deep排序,扫描哪些点对是符合的。 
	int l = 1,r = deep[0],sum = 0;
	while(l<r){
		if(deep[l]+deep[r]<=K){
			sum += r-l;
			++l;
		}
		else --r;
	}
	return sum;
}

void solve(int v){
	ans += cal(v,0);
	vis[v] = true;//标记根节点(相当于处理后,将根节点从子树中删除)。  
	for(int i=head[v] ; i ; i=E[i].next){
		if(!vis[E[i].to]){
			ans -= cal(E[i].to,E[i].value);//注意!这里要减去重复的。
			/*重复:比如有5个点a,b,c,d,e。其中a点为root。 
			有边a-b长为10,a-c长为10,c-d长为5,c-e长为5。 
			假如K为35,易知d到a和e到a的距离为15。
			因为15+15<35,则(d,e)满足条件。
			类似当root为c点后(d,e)也满足条件,此时就会出现重复。*/ 
			sum = son[E[i].to];
			root = 0;
			getroot(E[i].to,0);
			solve(root);//递归处理下一层。 
		}
	}
}

int main(){
	
	int a,b,c;
	while(scanf("%d %d",&N,&K) && (N || K)){
		memset(vis,false,sizeof vis);
		memset(head,0,sizeof head);
		top = ans = root = 0;
		for(int i=1 ; i<N ; i++){
			scanf("%d %d %d",&a,&b,&c);
			add_edge(a,b,c);
			add_edge(b,a,c);
		}
		f[0] = INF;
		sum = N;
		getroot(1,0);//先找到第一个重心。
		solve(root);
		printf("%d\n",ans);
	}
	
	return 0;
	
}