首页 > 其他分享 >ABC163F

ABC163F

时间:2022-11-11 21:36:16浏览次数:40  
标签:tmp cnt int siz sum ABC163F ans

正难则反,考虑用所有路径减去不包含颜色 \(k\) 的路径。

删掉所有颜色为 \(k\) 的节点,发现树分成了好多个连通块,设它们的大小为 \(s_1,s_2,\cdots\),那么不包含颜色 \(k\) 的路径条数就是

\[\sum \dfrac{s_i\times(s_i+1)}{2} \]

暴力做是 \(\mathcal O(n^2)\) 的。

设 \(f(x)=\dfrac{x\times(x+1)}{2}\)。

设 \(ans(i)\) 为颜色为 \(i\) 的答案,初始时均为 \(f(n)\)。

设 \(sum(i)\) 表示当前遍历过的点,颜色为 \(i\) 的节点的子树总大小,注意如果有包含就取最大的,\(siz(i)\) 表示以 \(i\) 为根的子树大小。

若当前在点 \(u\),先保存下此前的 \(sum(c_u)\),记为 \(tmp\),接下来要遍历它的某个儿子 \(v\),那么先清空 \(sum(c_u)\),然后进入 \(v\) 的子树,遍历完后 \(ans(c_u)\gets ans(c_u)-f(siz(v)-sum(c_u)),siz(u)\gets siz(u)+siz(v)\)。

当它的所有儿子遍历完后,更新 \(sum(c_u)\gets tmp+siz(u)\)。

注意不能漏下最上面的连通块,最后还要 \(ans(i)\gets ans(i)-f(n-sum(i))\)。

时间复杂度 \(\mathcal O(n)\)。

Code:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 200005;
int n;
int head[N], ver[N*2], nxt[N*2], cnt;
int c[N];
int sum[N], siz[N];
ll ans[N];

ll calc(int x) { return 1ll * x * (x + 1) / 2; }

void add(int u, int v) {
	ver[++cnt] = v, nxt[cnt] = head[u], head[u] = cnt;
}

void dfs(int u, int fa) {
	siz[u] = 1;
	int tmp = sum[c[u]];
	for (int i = head[u]; i; i = nxt[i]) {
		int v = ver[i];
		if (v == fa) continue;
		sum[c[u]] = 0, dfs(v, u), siz[u] += siz[v];
		ans[c[u]] -= calc(siz[v] - sum[c[u]]);
	}
	sum[c[u]] = tmp + siz[u];
}

int main() {
	scanf("%d", &n);
	for (int i = 1; i <= n; ++i) scanf("%d", &c[i]), ans[i] = calc(n);
	for (int i = 1, u, v; i < n; ++i) scanf("%d%d", &u, &v), add(u, v), add(v, u);
	dfs(1, 0);
	for (int i = 1; i <= n; ++i) ans[i] -= calc(n - sum[i]);
	for (int i = 1; i <= n; ++i) printf("%lld\n", ans[i]);
	return 0;
}

标签:tmp,cnt,int,siz,sum,ABC163F,ans
From: https://www.cnblogs.com/Kobe303/p/16882073.html

相关文章