小猪佩奇爬树

const int N=3e6+5;

int n,m,cnt,t[N],col[N],siz[N],Siz[N];
int gs[N],cs[N],head[N],Nod[N],Ans[N];

struct Node
{
    int nex,to;
};
Node e[N<<1];

inline void jia(int u,int v)
{
    e[++cnt].nex=head[u];
    head[u]=cnt;
    e[cnt].to=v;
}

inline void dfs(int u,int fa)
{
    siz[u]=1;
    int Col=col[u],sum=0;
    int οnend=0,Las=cs[Col];
    GO(i,u) 
    {
        int v=e[i].to;
        if(v==fa) continue;
        int las=cs[Col];//递归这颗子树之前的颜色个数
        dfs(v,u);
        Siz[u]+=siz[u]*siz[v];//对就1个该颜色答案的统计
        siz[u]+=siz[v];
        if(las^cs[Col]) sum++,οnend=v;//记录端点
    }
    Siz[u]+=siz[u]*(n-siz[u]);
    if(Las||t[Col]-1!=cs[Col]) sum++;//子树u自己本身也要统计进去
    cs[Col]++;
    if(sum==1) 
    {
        if(!gs[Col]) Nod[Col]=u;
        else 
        {
            int mul=(onend)?n-siz[onend]:siz[u];
            Ans[Col]=siz[Nod[Col]]*mul;
        }
        gs[Col]++;//统计有几个子树有颜色Col
    }
}

signed main()
{
    io.read(n);
    For(i,1,n)     
    {
        io.read(col[i]);
        t[col[i]]++;
        Nod[col[i]]=i;
    }
    For(i,1,n-1)
    {
        int x,y;
        io.read(x),io.read(y);
        jia(x,y),jia(y,x);
    }
    dfs(1,0);
    For(i,1,n) 
    {
        int ret=0;
        if(!t[i]) ret=n*(n-1)/2ll;//没有这种颜色
        if(t[i]==1) ret=Siz[Nod[i]]; //有1种
        if(gs[i]==2) ret=Ans[i];//两个端点情况
        io.write(ret),puts("");
    }
    return 0;
}