点分治板题

https://www.luogu.com.cn/problem/P4178
求一棵树上距离小于等于\(k\)的点对数，点分治解决

#include <cstdio>
#include <iostream>
#include <algorithm>
using namespace std;
const int N = 100000;
int n, k, mx[N], sz[N], siz, rt, la[N], ans, tot, dis[N], vis[N];
inline int read() {
	int res = 0, flag = 0; char ch = getchar();
	for(; !isdigit(ch); ch = getchar()) if(ch == '-') flag = 1;
	for(; isdigit(ch); ch = getchar()) res = (res << 1) + (res << 3) + (ch ^ 48);
	if(flag) res = ~res + 1;
	return res;
}
struct Edge{
	int to, nxt, w;
}e[100000];
inline int link(int u, int v, int w) {e[++tot] = (Edge) {v, la[u], w}, la[u] = tot;}
void getRoot(int u, int fa) {
	mx[u] = 0, sz[u] = 1;
	for(int i = la[u], v; i; i = e[i].nxt) {
		v = e[i].to;
		if(v == fa || vis[v]) continue;
		getRoot(v, u);
		mx[u] = max(mx[u], sz[v]), sz[u] += sz[v];
	}
	mx[u] = max(mx[u], siz - sz[u]);
	if(mx[u] < mx[rt]) rt = u;
}
void getDis(int u, int fa, int D) {
	dis[++dis[0]] = D;
	for(int v, i = la[u]; i; i = e[i].nxt) {
		v = e[i].to;
		if(vis[v] || v == fa) continue;
		getDis(v, u, D + e[i].w);
	}
	return ;
}
int calc(int u, int D) {
	int sum = 0;
	getDis(u, dis[0] = 0, D);
	sort(dis + 1, dis + dis[0] + 1);
	for(int l = 1, r = dis[0]; l <= r; ) {
		if(dis[l] + dis[r] <= k) sum += r - l, ++l;
		else --r;
	}
	return sum;
}
void solve(int u) {
	vis[u] = 1, ans += calc(u, 0);
	for(int i = la[u], v; i; i = e[i].nxt) {
		v = e[i].to;
		if(vis[v]) continue;
		ans -= calc(v, e[i].w);
		siz = sz[v], rt = 0, getRoot(v, 0), solve(v);
	}
}
int main() {
	mx[0] = 99999999;
	n = read();
	for(int i = 1, u, v, w; i < n; ++i) u = read(), v = read(), w = read(), link(u, v, w), link(v, u, w);
	k = read();
	siz = n, ans = 0; 
	getRoot(5, 0);
	solve(rt);
	printf("%d\n",ans);
	return 0;
}