POJ1741（点分治）

分治的时候SZ感觉是错的……但是貌似第一次找好重心就够了，之后SZ别太离谱就不会T，重心随一随缘就好……

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

const int maxn = 1e4 + 5;
int n, k, mx, SZ, ans;
struct Edge {
	int to, nxt, cost;
}e[maxn << 2];
int head[maxn], tot, size[maxn], root, dis[maxn], vis[maxn];
int l, r, Q[maxn];

void add(int u, int v, int c) {
	e[++tot].to = v, e[tot].cost = c, e[tot].nxt = head[u], head[u] = tot;
}

void GetRoot(int cur, int fa) {
	int tmp = 0;
	size[cur] = 1;
	for (int i = head[cur]; i; i = e[i].nxt) {
		int son = e[i].to;
		if (son == fa || vis[son])	continue;
		GetRoot(son, cur);
		size[cur] += size[son];
		tmp = max(tmp, size[son]);
	}
	tmp = max(tmp, SZ - size[cur]);
	if (tmp < mx)	mx = tmp, root = cur;
}

void GetDis(int cur, int fa) {
	Q[++r] = dis[cur];
	for (int i = head[cur]; i; i = e[i].nxt) {
		int son = e[i].to;
		if (son == fa || vis[son])	continue;
		dis[son] = dis[cur] + e[i].cost;
		GetDis(son, cur);
	}
}

int calc(int cur, int val) {
	l = 1, r = 0;
	dis[cur] = val;
	GetDis(cur, 0);
	sort(Q + 1, Q + 1 + r);
	int res = 0;
	while (l < r) {
		if (Q[l] + Q[r] <= k)	res += r - l, l++;
		else	r--;
	}
	return res;
}

void divide(int cur) {
	ans += calc(cur, 0);
	vis[cur] = 1;
	for (int i = head[cur]; i; i = e[i].nxt) {
		int son = e[i].to;
		if (vis[son])	continue;
		ans -= calc(son, e[i].cost);
		mx = 2e9, SZ = size[son];
		GetRoot(son, 0);
		divide(root);
	}
}

void init() {
	ans = 0, tot = 0, mx = 2e9, SZ = n;
	for (int i = 1; i <= n; i++)	head[i] = 0, vis[i] = 0;
}

int main() {
	while (scanf("%d %d", &n, &k) == 2 && (n | k)) {
		init();
		for (int u, v, cost, i = 1; i < n; i++) {
			scanf("%d %d %d", &u, &v, &cost);
			add(u, v, cost), add(v, u, cost);
		}
		GetRoot(1, 0);
		divide(root);
		printf("%d
", ans);
	}
	return 0;
}