poj 1741 Tree 树分治+重心

原题：http://poj.org/problem?id=1741

题意：

给定一个有N个点（编号0,1,…,N-1）的树，每条边都有一个权值（不超过1000）。

树上两个节点x与y之间的路径长度就是路径上各条边的权值之和。

求长度不超过K的路径有多少条。

 

对于任何边都可以分为两类：经过根节点和不经过根节点的，对于一棵树，可以先处理经过其根节点的，然后递归处理其子树，就变成了一个分治问题。

比如对于这么一棵树，我给点的编号1~6，边编号a~f

我们使用一个dis数组，dis[x]表示x到根节点的距离，那么位于两颗不同子树的点x，y的距离就等于dis[x]+dis[y]

假如dis[x]+dis[y]<=k  那么对于所有dis[i] <=dis[x] , dis[j]<=dis[y] , dis[i]+dis[j]<=k 都成立

那我们如何求dis呢，可以从根节点出发使用dfs构建dis数组，那么问题来了，对于两个点a,b，怎么知道他们是否属于两棵不同的子树呢？

要么可以另外使用一个数组记录该点所属的子树，要么可以直接不管，从最后的结果减去一些数字使得答案正确。

我选择的是第二个做法：

比如还是上面这颗子树，其不合法的点对（两个点都是同一棵子树）所连接的路径一定会经过根节点的左右儿子：dis[6]+dis[7]实际上是e+d+d+f这四条边的长度加起来。

那么我只需要减去  根节点的每个子儿子为根的子树符合条件的边的数量   即可。

 

另外，假如题目给的树是：

一条链，那么这个时候假如按1~n的顺序选根节点，时间复杂度就会退化为（N^2logN)

为了避免这种情况（实际上题目就卡了这种情况），我们还需要找出这棵树的重心，然后以该重心为根节点，时间复杂度就可以保持为（NlogNlogN)

重心的定义： 重心的最大子树的大小最小

详细可以看代码，我有注释

 

代码：

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
using namespace std;
const int INF = 0x3f3f3f3f;
const int maxn = 1e4 + 7;
int n, k, tot, hav_root, ans, id;
int Head[maxn], To[2 * maxn], Nxt[2 * maxn], Val[2 * maxn]; //这一堆是链式向前星的数组
int vis[maxn], siz[maxn], dis[maxn], d[maxn];
//vis ： 该点是否已经被删除  siz：子树大小  dis：距离根节点大小   d：用于将符合条件的dis放入然后排序用的
inline int read() { //快读
	int x = 0, f = 1; register char ch = getchar();
	for (; ch<'0' || ch>'9'; ch = getchar()) if (ch == '-') f = -1;
	for (; ch >= '0' && ch <= '9'; ch = getchar()) x = x * 10 + ch - '0';
	return x * f;
}
void init() {
	tot = 0;
	ans = 0;
	memset(Head, 0, sizeof(Head));
	memset(vis, 0, sizeof(vis));
}
inline void add_edge(int fro, int to, int val) { //链式向前星建边
	Nxt[++tot] = Head[fro];
	To[tot] = to;
	Val[tot] = val;
	Head[fro] = tot;
}
int min_siz;
void find_hav_root(int beg, int fat, int N) { // find havery root 找重心  beg:当前的点  fat：父亲节点  N：子树大小
	siz[beg] = 1;// 这里不需要vis 有fat就了
	int mx_part = 0;
	for (int i = Head[beg]; i; i = Nxt[i]) {
		int &to = To[i];
		if (vis[to] || to == fat) continue;
		find_hav_root(to, beg , N);
		siz[beg] += siz[to];
		mx_part = max(mx_part, siz[to]); //保存最大子树
	}
	mx_part = max(mx_part, N - siz[beg]);// 这里挺精髓的 ，因为父节点已经被vis过，所以父节点那一坨不会遍历，这样又可以减少时间
	if (min_siz > mx_part) {
		min_siz = mx_part;
		hav_root = beg;
	}
}
void get_dis(int x, int fat) { //得到dis数组  从x出发，父亲为fat
	d[++id] = dis[x]; //出来给人排序的
	for (int i = Head[x]; i; i = Nxt[i]) {
		int &to = To[i];
		if (vis[to] || to == fat) continue;
		dis[to] = dis[x] + Val[i];
		get_dis(to, x);
	}
}
int get_res(int x, int mit) { // 得到所有经过x点的边的数量   mit指的是x的初始距离 
	int res = 0;
	dis[x] = mit;
	id = 0; //id用于d数组
	get_dis(x, 0);
	sort(d + 1, d + 1 + id);
	int l = 1, r = id;
	while (l < r) {
		if (d[l] + d[r] <= k) res += (r - l++);
		else r--;
	}
	return res;
}
void solve(int root,int N) {
	vis[root] = 1;
	ans += get_res(root, 0);
	for (int i = Head[root]; i; i = Nxt[i]) {
		int &to = To[i];
		if (vis[to]) continue;
		ans -= get_res(to, Val[i]); //减去根节点的每个子儿子为根的子树符合条件的边
		min_siz = INF;
		int sto = siz[to]; // 子树搜的时候记得大小得变
		find_hav_root(to, 0, sto);
		solve(hav_root,sto);
	}
}
int main() {
	while (scanf("%d %d", &n, &k), n || k) {
		init();
		for (register int i = 1; i < n; i++) {
			int fro, to, val;
			fro = read();
			to = read();
			val = read();
			add_edge(fro, to, val); 
			add_edge(to, fro, val);
		}
		min_siz = INF;
		find_hav_root(1, 0, n);
		solve(hav_root,n);
		printf("%d\n", ans);
	}
	return 0;
}