POJ 1741 Tree

点分治

非常巧妙的实现, 复杂度 O(n * log ^2 n)
https://www.cnblogs.com/GXZlegend/p/6641720.html

#include <iostream>
#include <cstdio>
#include <cmath>
#include <algorithm>
#include <cstdlib>
#include <cstring>
using namespace std;
const int MAXN = 20005;
int head[MAXN], nume, n, k, siz[MAXN], masize, dep[MAXN];
int gra[MAXN], d[MAXN], tot, rot, ans;
bool f[MAXN];
int init() {
	int rv = 0, fh = 1;
	char c = getchar();
	while(c < '0' || c > '9') {
		if(c == '-') fh = -1;
		c = getchar();
	}
	while(c >= '0' && c <= '9') {
		rv = (rv<<1) + (rv<<3) + c - '0';
		c = getchar();
	}
	return fh * rv;
}
struct edge{
	int to, nxt, dis;
}e[MAXN];
void adde(int from, int to, int dis) {
	e[++nume].to = to;
	e[nume].nxt = head[from];
	e[nume].dis = dis;
	head[from] = nume;
}
void getroot(int u, int rt) {
	gra[u] = 0;
	siz[u] = 1;
	for(int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if(f[v] || v == rt) continue;
		getroot(v, u);
		siz[u] += siz[v];
		gra[u] = max(gra[u], siz[v]);
	}
	gra[u] = max(gra[u], masize - siz[u]);
	if(gra[rot] > gra[u]) rot = u;
}
void getdep(int u, int rt) {
	d[++tot] = dep[u];
	for(int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if(f[v] || v == rt) continue;
		dep[v] = dep[u] + e[i].dis;
		getdep(v, u);
	}
}
int calc(int u) {
	tot = 0;
	getdep(u, 0);
	sort(d + 1, d + 1 + tot);
	int l = 1, r = tot, rv = 0;
	while(l < r) {
		if(d[l] + d[r] <= k) rv += r - l, l++;
		else r--;
	}
	return rv;
}
void work(int u) {
	dep[u] = 0;
	f[u] = 1;
	ans += calc(u);
	for(int i = head[u]; i; i = e[i].nxt) {
		int v = e[i].to;
		if(!f[v]) {
			dep[v] = e[i].dis;
			ans -= calc(v);
			masize = siz[v];
			rot = 0;
			getroot(v, 0);
			work(rot);
		}
	}
}
int main() {
	while(1) {
		n = init(); k = init();
		if(!n || !k) break;
		memset(head, 0, sizeof(head));
		memset(f, 0, sizeof(f));
		nume = 0; tot = 0; masize = 0, ans = 0;
		for(int i = 1; i < n; i++) {
			int u = init(), v = init(), dis = init();
			adde(u, v, dis); adde(v, u, dis);
		}
		gra[0] = 0x3f3f3f3f;
		rot = 0;
		getroot(1, 0);
		work(rot);
		cout << ans << endl;
	}
	return 0;
}