C++ 解法, 执行用时: 2135ms, 内存消耗: 126804K, 提交时间: 2021-05-19 21:18:29

#include<bits/stdc++.h>
using namespace std;

vector <int> vec[1000005];
int n,a[1000005];
int siz[1000005],gg[1000005],del[1000005];
long long ans[1000005],add[1000005],dp[1000005],ggg[1000005];

void dfs0(int x,int fa)
{
	siz[x]=1;
	dp[x]=1;
	int tmp=gg[a[x]];
	gg[a[x]]=x;
	for(int i=0;i<vec[x].size();i++)
	{
		if(vec[x][i]!=fa)
		{
		    del[x]=0;
			dfs0(vec[x][i],x);
			siz[x]+=siz[vec[x][i]];
			add[vec[x][i]]=siz[vec[x][i]]-del[x];
			dp[x]+=dp[vec[x][i]]+add[vec[x][i]];
        }
	}
	del[tmp]+=siz[x];
	gg[a[x]]=tmp;
}

void dfs1(int x,int fa)
{
    ans[x]=n-siz[x]-ggg[a[x]];
    long long tmp=ggg[a[x]],now;
    ggg[a[x]]=n-siz[x]+1;
    for(int i=0;i<vec[x].size();i++)
	{
		if(vec[x][i]!=fa)
        {
            now=ggg[a[x]];
            dfs1(vec[x][i],x);
            ggg[a[x]]=now+siz[vec[x][i]];
		}
	}
	ggg[a[x]]=tmp+siz[x];
}
void dfs2(int x,int fa)
{
    ans[x]-=ggg[a[x]];
    ans[x]+=ans[fa]-add[x];
    long long tmp=ggg[a[x]],now;
    ggg[a[x]]=0;
    for(int i=vec[x].size()-1;i>=0;i--)
	{
		if(vec[x][i]!=fa)
        {
            now=ggg[a[x]];
            dfs2(vec[x][i],x);
            ggg[a[x]]=now+siz[vec[x][i]];
        }
	}
	ggg[a[x]]=tmp+siz[x];
}

int main()
{
	int u,v;
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&u,&v);
		vec[u].push_back(v);
		vec[v].push_back(u);
	}
	dfs0(1,0);
	ans[0]=dp[1];
	add[1]=0;
	dfs1(1,0);
	memset(ggg,0,sizeof(ggg));
	dfs2(1,0);
    for(int i=1;i<=n;i++)
		printf("%lld\n",ans[i]);
	return 0;
}