P5666 [CSP-S2019] 树的重心

P5666 CSP-S2019 树的重心

正难则反: 考虑一个点在删除哪些边的情况下会成为重心
将整棵树的重心作为根,设h[i]为i的重儿子的大小
性质: 删除i的子树中的边后i不会成为新连通块的重心
证明: 设之后i所在连通块大小为m 则剩下的sz'[i]=sz[i]-(n-m) 则m-sz'[i]=n-sz[i]>=n/2>m/2于是i不为重心
对于节点i,只需考虑不在i的子树中的边 ; 设树的大小减少了p
如果i为新的重心,则 n-p-sz[i]<=(n-p)/2 且 mx[i]<=(n-p)/2
于是 n-2sz[i] <= p <= n-2mx[i]
用树状数组维护p(修改的是i到其父节点的p)
1.初始化为sz[i] 2.到达时p变为n-sz[i] 3.离开时p变为sz[i]

注意: 要排除子树内p的贡献,遍历时用记录一个不恢复的新的树状数组
根节点的答案:设根的最大子树为x,次大子树为y
答案在x中: 2sz[y]<=n-sz[当前] 否则: 2sz[x]<=n-sz[当前]

点击查看代码

#include <stdio.h>
#include <string.h>
int max(int x, int y) { return x < y ? y : x; }
const int N = 3e5 + 5, M = N << 1;
typedef long long LL;
int T, n, h[N], e[M], nxt[M], idx;
int rt, sz[N], mx[N]; // 根,大小,最大大小
int c[N], d[N]; // 树状数组
void inc(int *tr, int x) { for(++ x; x <= n + 1; x += x & -x) ++ tr[x]; }
void dec(int *tr, int x) { for(++ x; x <= n + 1; x += x & -x) -- tr[x]; }
int query(int *tr, int x) {
	int y = 0;
	for(++ x; x; x -= x & -x) y += tr[x];
	return y;
}
bool in[N]; // 是否在u的子树中
int x, y;
LL res;
void add(int a, int b) {
	e[++ idx] = b, nxt[idx] = h[a], h[a] = idx;
}
void dfs1(int u, int fa) {
	sz[u] = 1, mx[u] = 0;
	for(int i = h[u], v; i; i = nxt[i]) {
		if((v = e[i]) == fa) continue;
		dfs1(v, u), sz[u] += sz[v], mx[u] = max(mx[u], sz[v]);
	}
}
void dfs2(int u, int fa) {
	int pre; // d数组原来的和
	if(fa) {
		dec(c, sz[u]), inc(c, n - sz[u]),
		res += LL(query(c, n - (mx[u] << 1)) - query(c, n - (sz[u] << 1) - 1)) * u,
			pre = query(d, n - (mx[u] << 1)) - query(d, n - (sz[u] << 1) - 1); // d数组原始值
		if(!in[u] && in[fa]) in[u] = true; // 是否在x的子树中
		if(sz[u] <= n - (sz[in[u] ? y : x] << 1)) res += rt; // 计算根的贡献
	}
	for(int i = h[u]; i; i = nxt[i]) if(e[i] != fa) dfs2(e[i], u);
	if(fa) {
		res -= LL(query(d, n - (mx[u] << 1)) - query(d, n - (sz[u] << 1) - 1) - pre) * u,
		inc(d, sz[u]), inc(c, sz[u]), dec(c, n - sz[u]);
	}
}
int main() {
	scanf("%d", &T);
	while(T --) {
		scanf("%d", &n);
		memset(h + 1, 0, n << 2), idx = 0, res = 0;
		for(int i = 1, a, b; i < n; i ++) scanf("%d%d", &a, &b), add(a, b), add(b, a);
		dfs1(1, 0); // 找出重心
		for(int i = 1; i <= n; i ++) if(max(mx[i], n - sz[i]) <= (n >> 1)) { rt = i; break; }
		dfs1(rt, 0); // 更新sz,mx
		memset(c + 1, 0, (n + 1) << 2), memset(d + 1, 0, (n + 1) << 2), memset(in + 1, 0, n);
		for(int i = 1; i <= n; i ++) if(i != rt) inc(c, sz[i]); // 初始化树状数组
		x = y = 0; // 找出最大和次大
		for(int i = h[rt], v; i; i = nxt[i])
			if(sz[v = e[i]] >= sz[x]) y = x, x = v;
			else if(sz[v] > sz[y]) y = v;
		in[x] = true;
		dfs2(rt, 0);
		printf("%lld\n", res);
	}
	return 0;
}