Poj1741/洛谷P4718 Tree（点分治）

题面

有多组数据：Poj

无多组数据：洛谷

题解

点分治板子题，$calc$的时候搞一个$two\ pointers$扫一下统计答案就行了。

#include <cmath>
#include <cstdio>
#include <cstring>
#include <algorithm>
using std::min; using std::max;
using std::swap; using std::sort;
typedef long long ll;

template<typename T>
void read(T &x) {
    int flag = 1; x = 0; char ch = getchar();
    while(ch < '0' || ch > '9') { if(ch == '-') flag = -flag; ch = getchar(); }
    while(ch >= '0' && ch <= '9') x = x * 10 + ch - '0', ch = getchar(); x *= flag;
}

const int N = 1e4 + 10, Inf = 1e9 + 7;
int n, k, d[N], Size;
int cnt, from[N], to[N << 1], nxt[N << 1], dis[N << 1];
int p, siz[N], tmp; bool vis[N];
inline void addEdge(int u, int v, int w) {
	to[++cnt] = v, nxt[cnt] = from[u], dis[cnt] = w, from[u] = cnt;
}

void getrt(int u, int f) {
	siz[u] = 1; int max_part = 0;
	for(int i = from[u]; i; i = nxt[i]) {
		int v = to[i]; if(v == f || vis[v]) continue;
		getrt(v, u), siz[u] += siz[v];
		max_part = max(max_part, siz[v]);
	} max_part = max(max_part, Size - siz[u]);
	if(max_part < tmp) tmp = max_part, p = u;
}

void getpoi(int x, int y, int f) {
	d[++tmp] = y;
	for(int i = from[x]; i; i = nxt[i]) {
		int v = to[i]; if(v == f || vis[v]) continue;
		getpoi(v, y + dis[i], x);
	}
}

int calc (int x, int y) {
	tmp = 0, getpoi(x, y, 0);
	sort(&d[1], &d[tmp + 1]);
	int j = tmp, ret = 0;
	for(int i = 1; i < j; ++i) {
		while(d[i] + d[j] > k && i < j) --j;
		ret += j - i;
	} return ret;
}

int doit(int x) {
	int ret = 0; 
	tmp = Inf, getrt(x, 0), vis[p] = true;
	ret += calc(p, 0);
	for(int i = from[p]; i; i = nxt[i]) {
		int v = to[i]; if(vis[v]) continue;
		ret -= calc(v, dis[i]), Size = siz[v], ret += doit(v);
	} return ret;
}

int main () {
	while(1) {
		read(n), read(k); if(!n && !k) break;
		Size = n, cnt = 0, memset(from, 0, sizeof from);
		memset(vis, 0, sizeof vis);
		for(int i = 1, u, v, w; i < n; ++i) {
			read(u), read(v), read(w);
			addEdge(u, v, w), addEdge(v, u, w);
		}
		printf("%d\n", doit(1));
	}
	return 0;
}