洛谷P4178：Tree（点分治模板）

题目大意：询问树上距离小于等于K的点对有多少个

#include<bits/stdc++.h>
using namespace std;
#define lowbit(i) (i & (-i))
const int maxn = 4e4 + 10;
int n,val[maxn + 10],k,root,f[maxn],sz[maxn],tot;
bool done[maxn];
#define pii pair<int,int>
#define fir first
#define sec second
vector<pii> g[maxn];
vector<int> t,tmp;
int ans = 0;
void upd(int x,int v) {
	for(; x <= maxn; x += lowbit(x)) val[x] += v;
}
int qry(int x) {
	int ans = 0;
	for(; x > 0; x -= lowbit(x)) ans += val[x];
	return ans;
}
void getroot(int u,int fa) {
	sz[u] = 1;f[u] = 0;
	for(auto it : g[u]) {
		if(done[it.fir] || it.fir == fa) continue;
		getroot(it.fir,u);
		sz[u] += sz[it.fir];
		f[u] = max(f[u],sz[it.fir]);
	} 
	f[u] = max(tot - sz[u],f[u]);
	if(f[u] < f[root] || !root) root = u;
}
void dfs(int u,int fa,int w) {
	t.push_back(w);tmp.push_back(w);
	for(auto it : g[u]) {
		if(done[it.fir] || it.fir == fa) continue;
		dfs(it.fir,u,w + it.sec);
	}
}
int solve(int u) {
	done[u] = true;
	int ans = 0;
	tmp.clear();
	for(auto it : g[u]) {
		if(done[it.fir]) continue;
		t.clear();
		dfs(it.fir,u,it.sec);
		for(auto i : t) {
			if(i > k) continue;
			ans += qry(k - i);
		}
		for(auto i : t) {
			if(i > k) continue;
			upd(i,1);
		}
	}
	ans += qry(k);
	for(auto i : tmp) {
		if(i > k) continue;
		upd(i,-1);
	}
	return ans;
}
void divide(int u) {
	ans += solve(u);
	for(auto it : g[u]) {
		if(done[it.fir]) continue;
		tot = sz[it.fir];root = 0;
		getroot(it.fir,-1);
		divide(it.fir);
	}
}
int main() {
	scanf("%d",&n);
	for(int i = 1; i < n; i++) {
		int u,v,w;
		scanf("%d%d%d",&u,&v,&w);
		g[u].push_back(pii(v,w));
		g[v].push_back(pii(u,w));
	}
	scanf("%d",&k);
	tot = n;root = 0;
	getroot(1,-1);
	divide(1);
	printf("%d\n",ans);
	return 0;
}