BZOJ1468 点分治入门练习题

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
using namespace std;
const int maxn = 100007, inf = ~0u >> 1;
int n, k, c, csiz, tot, size[maxn], dis[maxn], arr[maxn];
int cnt, head[maxn], next[maxn << 1], to[maxn << 1], len[maxn << 1];

void addedge(int u, int v, int l) {
	next[++cnt] = head[u]; to[cnt] = v;
	len[cnt] = l; head[u] = cnt;
}
bool vis[maxn];
void DFS(int u, int last) {
	size[u] = 1; int mxsz = 0;
	for(int i = head[u]; i; i = next[i]) {
		if(to[i] == last || vis[to[i]]) continue;
		DFS(to[i], u);
		size[u] += size[to[i]];
		mxsz = max(mxsz, size[to[i]]);
	}
	mxsz = max(mxsz, tot - size[u]);
	if(csiz > mxsz) csiz = mxsz, c = u;
}

void Dist(int u, int last) {
	arr[++arr[0]] = dis[u];
	for(int i = head[u]; i; i = next[i]) {
		if(vis[to[i]] || to[i] == last) continue;
		dis[to[i]] = dis[u] + len[i];
		Dist(to[i], u);
	}
}

int Calc(int u, int l) {
	dis[u] = l; arr[0] = 0;
	Dist(u, 0);
	sort(arr + 1, arr + arr[0] + 1);
	int ret = 0;
	for(int i = 1, j = arr[0]; i < j;) {
		if(arr[i] + arr[j] > k) j--;
		else ret += j - i, i++;
	} 
	return ret;
}

int ret;
void get_ans(int u) {
	tot = size[u] ? size[u] : n;
	csiz = inf; DFS(u, 0); vis[u = c] = 1;
	ret += Calc(u, 0);
	for(int i = head[u]; i; i = next[i]) {
		if(vis[to[i]]) continue;
		ret -= Calc(to[i], len[i]);
		get_ans(to[i]);
	}
}

int main() {
	scanf("%d", &n);
	for(int i = 1; i < n; i++) {
		int u, v, l; scanf("%d%d%d", &u, &v, &l);
		addedge(u, v, l); addedge(v, u, l);
	}
	scanf("%d", &k);
	get_ans(1);
	printf("%d\n", ret);
	return 0;
}