POJ1741--Tree(树的分治)

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

题意：求树上的满足dis(u,v) <= k 的点对数

思路：首先要理解什么是的重心。不知道的先切POJ1655.

      每次找到树的重心，那么满足条件的u,v可以分为2种情况，经过和不经过重心。找到重心后，求bfs得到可达点的距离（之前的重心视为不可达）。然后推入数组后排序。维护2个指针，一个从头到尾，另一个从尾到头。得到dis(u,v)<=k的个数，这样会算多。得减去子树中dis(u,v)<=k的个数。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>
using namespace std;
#define maxn 20080
#define inf 0x3f3f3f3f
#define LL long long int
int first[maxn];
int vv[maxn],nxt[maxn],ww[maxn],cnt[maxn],dp[maxn],dis[maxn],key[maxn];
bool vis[maxn];
int e,n,Center,Size,K;
int cc[maxn],dd[maxn],num1,num2,num;
LL ans;
inline int max(int a,int b)
{
	return a>b?a:b;
}

void init()
{
	e = ans = 0;
	memset(first,-1,sizeof(first));
	memset(vis,0,sizeof(vis));
}

void addedge(int u,int v,int w)
{
	vv[e] = v;	ww[e] = w;	nxt[e] = first[u];	first[u] = e++;
	vv[e] = u;	ww[e] = w;	nxt[e] = first[v];	first[v] = e++;
}

void dfs(int u,int fa)			///先dfs求出整棵树的大小。每个节点子树的大小，每个节点最大的子树
{
	num++;
	cnt[u] = 1;
	dp[u] = 0;
	for(int i = first[u];i != -1;i = nxt[i])
	{
		int v = vv[i];
		if(v == fa || vis[v])	continue;
		dfs(v,u);
		cnt[u] += cnt[v];
		dp[u] = max(cnt[v],dp[u]);
	}
}

void get_center(int u,int fa)	///得到树的重心
{
	int tmp = max(num-cnt[u],dp[u]);
	if(tmp < Size)
	{
		Center = u;
		Size = tmp;
	}
	for(int i = first[u];i != -1;i = nxt[i])
	{
		int v = vv[i];
		if(v == fa || vis[v])	continue;
		get_center(v,u);
	}
}

void dfs2(int u,int fa)		///搜树，将距离搞进去
{
	if(dis[u] > K)	return;
	dd[num2++] = dis[u];
	cc[num1++] = dis[u];
	for(int i = first[u];i != -1;i = nxt[i])
	{
		int v = vv[i];
		if(vis[v] || v == fa)	continue;
		dis[v] = dis[u] + ww[i];
		dfs2(v,u);
	}
}

void Get_Center(int u,int pre)//得到树的重心
{
	Center = u,Size = n;
	dfs(u,pre);
	get_center(u,pre);
}

void DFS(int u,int pre)	//这个深搜用来得到答案。
{
	num1 = num2 = num = 0;
	cc[num1++] = 0;
	Get_Center(u,pre);	//得到树的重心
	vis[Center] = 1;
	dis[Center] = 0;
	for(int i = first[Center];i != -1;i = nxt[i])
	{
		num2 = 0;
		int v = vv[i];
		if(!vis[v] && v != pre)	
		{
			dis[v] = ww[i];
			dfs2(v,Center);
			sort(dd,dd+num2);
			int tou = 0,wei = num2-1;
			while(tou < wei)
			{
				if(dd[tou] + dd[wei] <= K)	ans -= wei - tou;
				else 
				{
					while(dd[tou] + dd[wei] > K)
					{
						wei--;
						if(wei <= tou)	break;
					}
					ans -= wei - tou;
				}
				tou++;
			}
		}
	}
	sort(cc,cc+num1);
	int tou = 0,wei = num1-1;
	while(tou < wei)
	{
		if(cc[tou] + cc[wei] <= K)	ans += wei - tou;
		else 
		{
			while(cc[tou] + cc[wei] > K)
			{
				wei--;
				if(wei <= tou)	break;
			}
			ans += wei - tou;
		}
		tou++;
	}
	for(int i = first[Center];i != -1;i = nxt[i])
	{
		int v = vv[i];
		if(!vis[v])	DFS(v,Center);
	}
}

int main()
{
	//freopen("in.txt","r",stdin);
	while(scanf("%d%d",&n,&K)==2 && n|K)
	{
		init();
		for(int i = 1;i < n;i++)
		{
			int u,v,w;
			scanf("%d%d%d",&u,&v,&w);
			addedge(u,v,w);
		}
		DFS(1,0);
		printf("%lld\n",ans);
	}
	return 0;
}