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;
}