Tree POJ - 1741 （点分治模板）

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

Sample Output

题意：

求树上两点间距离小于等于k的点的对数。

大致的解释写代码里了，理解透了再详细写一遍。。。（/逃）

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

typedef long long ll;
const int maxn = 1e4 + 100;

int top, num, root, Max, ans, n, k, cnt;
int head[maxn], size[maxn], dist[maxn];
bool vis[maxn];

struct node {
	int v, w, next;
}edge[maxn << 1];

inline void add(int u, int v, int w) {     //链式前向星存树
	edge[top].v = v;
	edge[top].w = w;
	edge[top].next = head[u];
	head[u] = top++;
}

inline void Init() {            //初始化
	top = 0;
	ans = 0;
	memset(vis, false, sizeof(vis));
	memset(head, -1, sizeof(head));
}

void dfs_root(int u, int father) {     //用来找重心
	size[u] = 1;      //size数组记录以节点i为根的子树的节点数
	int ma = 0;
	for(int i = head[u]; i != -1; i = edge[i].next) {
		int v = edge[i].v, w = edge[i].w;
		if(v != father && !vis[v]) {
			dfs_root(v, u);          
			size[u] += size[v];     
			ma = max(ma, size[v]); 
		}
	}
	ma = max(ma, cnt - size[u]);
	if(ma < Max) {               //通过比较得出重心
		Max = ma;
		root = u;
	}
}

void dfs_dis(int u, int father, int dis) {    //求以某个节点为根时，其子树所有节点到根的距离
	dist[num++] = dis;      //将所有距离存放到dist数组中
	for(int i = head[u]; i != -1; i = edge[i].next) {
		int v = edge[i].v, w = edge[i].w;
		if(v != father && !vis[v]) {
			dfs_dis(v, u, dis + w);
		}
	}
}

int calc(int u, int dis) {    //计算相应子树的答案
	int res = 0;
	num = 0;
	dfs_dis(u, 0, dis);       //求u的子树中所有节点到u的距离
	sort(dist, dist + num);   //排序是为了方便以O(n)的复杂度解决计数问题
	int i = 0, j = num - 1;
	while(i < j) {
		while(i < j && dist[i] + dist[j] > k)
			j --;
		res += j - i;
		i ++;
	}
	return res;
}

void work(int u)
{
	Max = n;     //用来求重心
	dfs_root(u, 0);    //找到当前子树的重心
	cnt = size[u];
	ans += calc(root, 0);    //计数
	vis[root] = true;        //分治过的节点打上标记
	for(int i = head[root]; i != -1; i = edge[i].next) {
		int v = edge[i].v, w = edge[i].w;
		if(!vis[v]) {
			ans -= calc(v, w);   //减去不合法的
			work(v);
		}
	}
}

int main()
{
	//freopen("in.txt", "r", stdin);
	while(scanf("%d%d", &n, &k) != EOF) {
		if(n == 0 && k == 0)	break;
		int u, v, w;
		Init();
		for(int i = 1; i <= n - 1; ++ i) {
			scanf("%d%d%d", &u, &v, &w);
			add(u, v, w);
			add(v, u, w);
		}
		work(1);
		cout << ans << endl;
	}
	return 0;
}

Tree POJ - 1741 （点分治模板）

猜你喜欢