【POJ1741】Tree【点分治】

题意：

n个点的树，求距离不大于k的点对个数。

点分治模板题。

#include <cstdio>
#include <vector>
#include <algorithm>

using namespace std;

const int maxn = 10005, inf = 0x3f3f3f3f;

int n, k, head[maxn], cnt, dis[maxn], root, ans, tot, size[maxn], f[maxn];
bool vis[maxn];

vector<int> depth;

struct _edge {
	int v, w, next;
} g[maxn << 1];

inline int iread() {
	int f = 1, x = 0; char ch = getchar();
	for(; ch < '0' || ch > '9'; ch = getchar()) f = ch == '-' ? -1 : 1;
	for(; ch >= '0' && ch <= '9'; ch = getchar()) x = x * 10 + ch - '0';
	return f * x;
}

inline void add(int u, int v, int w) {
	g[cnt] = (_edge) {v, w, head[u]};
	head[u] = cnt++;
}

void getroot(int x, int p) {
	size[x] = 1; f[x] = 0;
	for(int i = head[x]; ~i; i = g[i].next) if(g[i].v ^ p && !vis[g[i].v]) {
		getroot(g[i].v, x);
		size[x] += size[g[i].v];
		f[x] = max(f[x], size[g[i].v]);
	}
	f[x] = max(f[x], tot - size[x]);
	if(f[x] <= f[root]) root = x;
}

void getdepth(int x, int p) {
	depth.push_back(dis[x]);
	size[x] = 1;
	for(int i = head[x]; ~i; i = g[i].next) if(g[i].v ^ p && !vis[g[i].v]) {
		dis[g[i].v] = dis[x] + g[i].w;
		getdepth(g[i].v, x);
		size[x] += size[g[i].v];
	}
}

inline int calc(int x, int d) {
	depth.clear(); dis[x] = d;
	getdepth(x, 0);
	sort(depth.begin(), depth.end());
	int l = 0, r = depth.size() - 1, res = 0;
	while(l < r)
		if(depth[l] + depth[r] <= k) res += r - l++;
		else r--;
	return res;
}

void work(int x) {
	ans += calc(x, 0);
	vis[x] = 1;
	for(int i = head[x]; ~i; i = g[i].next) if(!vis[g[i].v]) {
		ans -= calc(g[i].v, g[i].w);
		tot = f[root = 0] = size[g[i].v];
		getroot(g[i].v, 0);
		work(root);
	}
}

int main() {
	while(1) {
		n = iread(); k = iread();
		if(n == 0 && k == 0) break;

		for(int i = 1; i <= n; i++) head[i] = -1, vis[i] = 0; cnt = 0;

		for(int i = 1; i < n; i++) {
			int u = iread(), v = iread(), w = iread();
			add(u, v, w); add(v, u, w);
		}

		tot = f[root = 0] = n;
		getroot(1, 0);
		ans = 0;
		work(root);
		printf("%d\n", ans);
	}
	return 0;
}