点分治板题
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;
}