2023.9.3 P2664

P2664 树上游戏

对于颜色 \(c\)，如果我们把颜色 \(c\) 的点全部都删除，那么我们会得到若干个连通块。
连通块里面的路径是没有贡献的，连通块联通外面的路径都会有这个颜色做了贡献。
对于一个连通块，其里面所有点都能有 \(n-siz(连通块)\) 的贡献。
如果我们每次枚举颜色，再计算连通块，这样是 \(O(n^2)\).
那么我们考虑一次把所有颜色计算出来。

若当前 dfs 到一个 \(u\)，其颜色是 \(c\)，那么删除它，其所有儿子 \(v_1,v_2,...\) 都会各自形成一个连通块。
遍历其儿子。我们现在要求出这个儿子 \(v\) 所在连通块的大小。
如何求呢？这个连通块的大小是 \(v\) 的子树大小减去 \(v\) 子树内极大的、颜色也是 \(c\) 的子树大小。
用一个 vector 存储现在颜色是 \(c\) 的极大子树有哪些，和一个 \(cs\) 数组表示颜色为 \(c\) 的子树大小之和。
我们在 dfs 到 \(v\) 前后计算 \(cs_c\) 的增量 \(del\)， 其代表的是 \(v\) 子树内颜色为 \(c\) 的极大子树大小的和。
那么这个连通块的大小就是 \(siz_v-del\) 了。
树上差分，\(v\) 的连通块的所有点贡献都加上了 \(n-(siz_v-del)\)，
具体是在 \(v\) 加上贡献，在 \(v\) 内的极大子树减去贡献，然后从根节点开始加这个差分。

code

#include<bits/stdc++.h>
using namespace std;
const int N=2e5+10;
int n,m,col[N],dfn[N],num,siz[N],cs[N];
long long d[N];
vector<int> e[N];
vector<int> vec[N];
void dfs(int u,int fa) {
	dfn[u]=++num; siz[u]=1;
	int c=col[u];
	for(int v:e[u]) {
		if(v==fa) continue;
		int cs0=cs[c];
		dfs(v,u);
		int ss=siz[v]-(cs[c]-cs0);
		cs[c]=siz[v]+cs0;
		d[v]+=ss;
		siz[u]+=siz[v];
		for(; vec[c].size(); ) {
			if(dfn[vec[c].back()]>dfn[u]) {
				d[vec[c].back()]-=ss;
				vec[c].pop_back();
			} else break;
		}
	}
	vec[c].push_back(u);
	cs[c]++;
}
void dfs2(int u,int fa) {
	for(int v:e[u]) {
		if(v==fa) continue;
		d[v]+=d[u];
		dfs2(v,u);
	}
}
int main() {
	scanf("%d",&n);
	for(int i=1; i<=n; i++) scanf("%d",&col[i]),m=max(m,col[i]);
	for(int i=1,u,v; i<n; i++) {
		scanf("%d%d",&u,&v);
		e[u].push_back(v); e[v].push_back(u);
	}
	dfs(1,1);
	int tot=0;
	for(int c=1; c<=m; c++) {
		if(vec[c].size()) {
			tot++;
			d[1]+=n-cs[c];
			for(; vec[c].size(); ) {
				d[vec[c].back()]-=n-cs[c];
				vec[c].pop_back();
			}
		}
	}
	dfs2(1,1);
	for(int i=1; i<=n; i++) {
		printf("%lld\n",1ll*tot*n-d[i]);
	}
	return 0;
}