POJ 1741 Tree 树分治

题意：

给出一颗有\(n (n \leq 10^4)\)个节点的树，和一个\(k\)。统计有多少个点对\(u, \, v(u \neq v)\)满足\(u\)到\(v\)的最短距离不超过\(k\)。

分析：

树分治的入门题，可以参考论文《分治算法在树的路径问题中的应用》。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <map>
using namespace std;

typedef pair<int, int> PII;

const int maxn = 10000 + 10;
const int INF = 0x3f3f3f3f;

struct Edge
{
	int v, w, nxt;
	Edge() {}
	Edge(int v, int w, int nxt): v(v), w(w), nxt(nxt) {}
};

int n, k;

int ecnt, head[maxn];
Edge edges[maxn * 2];

void AddEdge(int u, int v, int w) {
	edges[ecnt] = Edge(v, w, head[u]);
	head[u] = ecnt++;
}

int sz[maxn], fa[maxn];  //sz[u]是以u为根子树的大小，fa[u]是u的父亲
bool del[maxn];  //del[u]表示u作为某颗子树的重心删除的标记
int ans, mins, centroid;  //最终答案，最小的最大子树以及重心
vector<int> d, d2;

//计算sz和fa
void dfs(int u) {
	sz[u] = 1;
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(v == fa[u] || del[v]) continue;
		fa[v] = u;
		dfs(v);
		sz[u] += sz[v];
	}
}

//计算重心
void dfs2(int u, int t) {
	int m = 0;
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(v == fa[u] || del[v]) continue;
		m = max(m, sz[v]);
		dfs2(v, t);
	}
	m = max(m, t - sz[u]);
	if(m < mins) { mins = m; centroid = u; }
}

//统计所有点到根节点的距离
void getdist(int u, int p, int dist, vector<int>& d) {
	d.push_back(dist);
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v, w = edges[i].w;
		if(v == p || del[v]) continue;
		getdist(v, u, dist + w, d);
	}
}

//统计符合要求点对个数
int cntpair(vector<int>& d) {
	int ans = 0;
	sort(d.begin(), d.end());
	int j = d.size();
	for(int i = 0; i < d.size(); i++) {
		while(j > 0 && d[i] + d[j-1] > k) j--;
		ans += j - (j > i ? 1 : 0);  //去掉和自己成为一对的情况
	}
	return ans / 2;
}

//分治过程
void solve(int u) {
	fa[u] = 0;
	dfs(u);
	mins = INF;
	dfs2(u, sz[u]);
	int s = centroid;
	del[s] = true;

	for(int i = head[s]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(del[v]) continue;
		solve(v);
	}

	d.clear();
	d.push_back(0);
	for(int i = head[s]; ~i; i = edges[i].nxt) {
		int v = edges[i].v, w = edges[i].w;
		if(del[v]) continue;
		d2.clear();
		getdist(v, s, w, d2);
		ans -= cntpair(d2);  //去掉计重部分
		d.insert(d.end(), d2.begin(), d2.end());
	}
	ans += cntpair(d);
	del[s] = false;
}

int main()
{
	while(scanf("%d%d", &n, &k) == 2) {
		if(!n && !k) break;
		ecnt = 0;
		memset(head, -1, sizeof(head));
		for(int i = 1; i < n; i++) {
			int u, v, w; scanf("%d%d%d", &u, &v, &w);
			AddEdge(u, v, w);
			AddEdge(v, u, w);
		}
		
		ans = 0;
		solve(1);
		printf("%d\n", ans);
	}

	return 0;
}