某模拟赛C题 树上路径统计 （点分治）

题意

给定一棵有n个节点的无根树，树上的每个点有一个非负整数点权。定义一条路径的价值为路径上的点权和-路径上的点权最大值。 给定参数P，我！=们想知道，有多少不同的树上简单路径，满足它的价值恰好是P的倍数。 注意：单点算作一条路径；u!=v时，(u,v)和（v,u）只算一次。

题解

树上路径统计，解法是点分治。点分的时候求出根到每个点路径最大值和权值和。排一序，然后开个桶，就能计算了。去重就套路的减去没棵子树里面的答案。

CODE

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100005;
typedef long long LL;
LL ans;
int n, mod, fir[MAXN], nxt[MAXN<<1], to[MAXN<<1], cnt, val[MAXN];
inline void link(int x, int y) {
	to[++cnt] = y; nxt[cnt] = fir[x]; fir[x] = cnt;
	to[++cnt] = x; nxt[cnt] = fir[y]; fir[y] = cnt;
}
bool ban[MAXN];
int getsz(int u, int ff) {
	int re = 1;
	for(int v, i = fir[u]; i; i = nxt[i])
		if((v=to[i]) != ff && !ban[v])
			re += getsz(v, u);
	return re;
}
int getrt(int u, int ff, int &rt, int Size) {
	int re = 1; bool can = 1;
	for(int v, tmp, i = fir[u]; i; i = nxt[i])
		if((v=to[i]) != ff && !ban[v]) {
			re += (tmp = getrt(v, u, rt, Size));
			if((tmp<<1) > Size) can = 0;
		}
	if(((Size-re)<<1) > Size) can = 0;
	if(can) rt = u;
	return re;
}
struct node {
	int mx, v;
	inline bool operator <(const node &o)const {
		return mx < o.mx;
	}
}seq[MAXN], vv[MAXN];
int tot;
void dfs(int u, int ff, int mx, int vs) {
	vs = (vs + val[u]) % mod;
	mx = max(mx, val[u]);
	vv[u] = (node){ mx, vs };
	for(int v, i = fir[u]; i; i = nxt[i])
		if((v=to[i]) != ff && !ban[v])
			dfs(v, u, mx, vs);
}
void push(int u, int ff) {
	seq[++tot] = vv[u];
	for(int v, i = fir[u]; i; i = nxt[i])
		if((v=to[i]) != ff && !ban[v])
			push(v, u);
}
int bin[10000005];
LL calc(int rt, int o) {
	tot = 0; push(rt, 0);
	sort(seq + 1, seq + tot + 1);
	LL re = 0;
	for(int i = 1; i <= tot; ++i) {
		re += bin[((seq[i].mx+o-seq[i].v)%mod+mod)%mod];
		++bin[seq[i].v%mod];
	}
	for(int i = 1; i <= tot; ++i) --bin[seq[i].v%mod];
	return re;
}
void solve(int x) {
	dfs(x, 0, 0, 0);
	ans += calc(x, val[x]);
	ban[x] = 1;
	for(int v, i = fir[x]; i; i = nxt[i])
		if(!ban[v=to[i]]) ans -= calc(v, val[x]);
}
void TDC(int x) {
	int Size = getsz(x, 0);
	getrt(x, 0, x, Size);
	solve(x);
	for(int v, i = fir[x]; i; i = nxt[i])
		if(!ban[v=to[i]]) TDC(v);
}
int main () {
	scanf("%d%d", &n, &mod);
	for(int i = 1, x, y; i < n; ++i)
		scanf("%d%d", &x, &y), link(x, y);
	for(int i = 1; i <= n; ++i) scanf("%d", &val[i]);
	TDC(1);
	printf("%lld\n", ans+n);
}