luogu P2664 树上游戏

题面传送门
这种数颜色的东西考虑颜色分开考虑。
如果我们将这些颜色对应的点删去，那么对于每个点对应的联通块大小\(siz\),这个颜色对这个点的贡献就是\(n-siz\)
所以我们需要算的就是每个点的联通块大小。
可以记\(P_i\)为\(i\)的父亲对应的颜色中在\(i\)子树内且与\(i\)间没有其它相同颜色的节点的子树大小和。可以发现在\(i\)这个点需要做的有些事情：
要让这个点子树内加上\(siz_i-P_i\)
要让这个点对应颜色撤销祖先的推下来的标记。
这个用一个栈预处理一下即可。
时间复杂度\(O(n)\)
code:

#include<bits/stdc++.h>
#define I inline
#define max(a,b) ((a)>(b)?(a):(b))
#define min(a,b) ((a)<(b)?(a):(b))
#define abs(x) ((x)>0?(x):-(x))
#define re register
#define RI re int
#define ll long long
#define db double
#define lb long db
#define N (100000+5)
#define M (200000)
#define mod 1000000007
#define Mod (mod-1)
#define eps (1e-9)
#define U unsigned int
#define it iterator
#define Gc() getchar() 
#define Me(x,y) memset(x,y,sizeof(x))
#define Mc(x,y) memcpy(x,y,sizeof(x))
#define d(x,y) (n*(x-1)+(y))
#define R(n) (rand()*rand()%(n)+1)
#define Pc(x) putchar(x)
#define LB lower_bound
#define UB upper_bound
#define PB push_back
using namespace std;vector<int> S[N+5];
int n,m,k,x,y,Maxn,C[N+5],Si[N+5],F[N+5],Sf[N+5],cnt,Fa[N+5];ll Ans[N+5],W[N+5],P[N+5],D[N+5];
I void Make(int x,int La){Si[x]=1;for(RI i:S[x]) i^La&&(Make(i,x),Si[x]+=Si[i]);}
I void dfs1(int x,int La){int Lp=F[C[La]];x^1&&(F[C[La]]=x);Fa[x]=F[C[x]];Fa[x]^1?(P[Fa[x]]+=Si[x]):(D[C[x]]+=Si[x]);for(RI i:S[x]) i^La&&(dfs1(i,x),0);F[C[La]]=Lp;}
I void dfs(int x,int La){
	if(x==1){Ans[x]+=1ll*Si[x]*cnt;for(RI i=1;i<=Maxn;i++) Ans[x]-=D[i];}else Ans[x]+=Si[x]-P[x],Ans[x]-=(Fa[x]^1?Si[Fa[x]]-P[Fa[x]]:Si[1]-D[C[x]]);
	for(RI i:S[x])i^La&&(dfs(i,x),0);
}
I void calc(int x,int La){Ans[x]+=Ans[La];for(RI i:S[x]) i^La&&(calc(i,x),0);}
int main(){
	freopen("1.in","r",stdin);
	RI i;scanf("%d",&n);for(i=1;i<=n;i++) scanf("%d",&C[i]),Maxn=max(Maxn,C[i]),cnt+=!Sf[C[i]],Sf[C[i]]=1,F[C[i]]=1;for(i=1;i<n;i++) scanf("%d%d",&x,&y),S[x].PB(y),S[y].PB(x);Make(1,0);
	dfs1(1,0);dfs(1,0);
	calc(1,0);for(i=1;i<=n;i++) printf("%lld\n",1ll*n*cnt-Ans[i]);
}