树上游戏

#include <iostream>
#include <cstdio>
#define LL long long
using namespace std;
int read()
{
	int x=0,f=1;
	char c=getchar();
	while(c<'0'||c>'9')
	{
		if(c=='-')  f=-1;
		c=getchar();
	}
	while(c>='0'&&c<='9')
	{
		x=x*10+c-'0';
		c=getchar();
	}
	return f*x;
}
const int N=1e5+5;
int root,size[N],sum,maxson[N],vis[N];
int del[N],cnt,rem[N],num,qua[N];
int n,a1,b1,co[N],had[N],to[N*2],nxt[N*2],p;
LL ans[N],val[N],wei[N],tot,s;
bool f[N];
void add_edge(int x,int y)
{
	nxt[++p]=had[x];
	to[p]=y;
	had[x]=p;
}
void getroot(int u,int fa)
{
	size[u]=1;  maxson[u]=0;
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(v==fa||f[v])  continue;
		getroot(v,u);
		size[u]+=size[v];
		maxson[u]=max(maxson[u],size[v]);
	}
	maxson[u]=max(maxson[u],sum-size[u]);
	if(maxson[u]<maxson[root])  root=u;
}
void getsize(int u,int fa)
{
	size[u]=1;
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(v==fa||f[v])  continue;
		getsize(v,u);
		size[u]+=size[v];
	}
}
void gettot(int u,int fa)
{
	wei[u]=0;
	if(!vis[co[u]])
	{
		rem[++num]=co[u];
		tot+=size[u];
		val[co[u]]+=size[u];
		wei[u]+=size[u];
	}
	vis[co[u]]++;
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(v==fa||f[v])  continue;
		gettot(v,u);
		wei[u]+=wei[v];
	}
	vis[co[u]]--;
}
void getqua(int u,int fa)
{
	if(!vis[co[u]])
	{
		del[++cnt]=co[u];
		qua[co[u]]+=size[u];
	}
	vis[co[u]]++;
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(v==fa||f[v])  continue;
		getqua(v,u);
	}
	vis[co[u]]--;
}
void getans(int u,int fa,LL mul,LL add)
{
	if(!vis[co[u]])
	{
		add-=val[co[u]]-qua[co[u]];
		mul++;
	}
	ans[u]+=mul*s+add;
	vis[co[u]]++;
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(v==fa||f[v])  continue;
		getans(v,u,mul,add);
	}
	vis[co[u]]--;
}
void solve(int u)
{
	num=0;  tot=0;  vis[co[u]]++;
	getsize(u,0);
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(f[v])  continue;
		gettot(v,u);
	}
	ans[u]+=tot+sum-1;
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(f[v])  continue;
		cnt=0;  s=sum-size[v];
		getqua(v,u);
		getans(v,u,1,tot-wei[v]);
		for(int j=1;j<=cnt;j++)  qua[del[j]]=0;
	}
	for(int i=1;i<=num;i++)  val[rem[i]]=0;
	vis[co[u]]--;
}
void divide(int u)
{
	f[u]=1;  solve(u);
	for(int i=had[u];i;i=nxt[i])
	{
		int v=to[i];
		if(f[v])  continue;
		root=0;  sum=size[v];
		getroot(v,0);  divide(root);
	}
}
int main()
{
	n=read();
	for(int i=1;i<=n;i++)  co[i]=read();
	for(int i=1;i<n;i++)
	{
		a1=read();  b1=read();
		add_edge(a1,b1);  add_edge(b1,a1);
	}
	maxson[root]=sum=n;
	getroot(1,0);  divide(root);
	for(int i=1;i<=n;i++)  cout<<ans[i]+1<<"\n";
	return 0;
}