洛谷P2664 树上游戏

https://www.luogu.org/problemnew/show/P2664

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<vector>
using namespace std;
#define fi first
#define se second
#define mp make_pair
#define pb push_back
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
struct E
{
    int to,nxt;
}e[200111];
int f1[100011],ne;
int sz[100011],a[100011];
int n;
ll t1[100011],t2[100011],s,ans[100011];
void dfs1(int u,int fa)
{
    sz[u]=1;
    int v;
    ll t=t1[a[u]],z=t1[a[fa]];
    for(int k=f1[u];k;k=e[k].nxt)
        if(e[k].to!=fa)
        {
            v=e[k].to;
            dfs1(v,u);
            sz[u]+=sz[v];
        }
    t1[a[u]]=t+sz[u];
    t2[u]=t1[a[fa]]-z;
}
void dfs2(int u,int fa)
{
    int v;ll ta;
    ans[u]=s;
    for(int k=f1[u];k;k=e[k].nxt)
        if(e[k].to!=fa)
        {
            v=e[k].to;
            ta=t1[a[v]];
            s+=n-t1[a[v]];
            t1[a[v]]=n;
            s+=t2[v]-sz[v];
            t1[a[u]]+=t2[v]-sz[v];
            dfs2(v,u);
            s+=ta-t1[a[v]];
            t1[a[v]]=ta;
            s-=t2[v]-sz[v];
            t1[a[u]]-=t2[v]-sz[v];
        }
}
int main()
{
    int i,x,y;
    scanf("%d",&n);
    for(i=1;i<=n;++i)
        scanf("%d",a+i);
    for(i=1;i<n;++i)
    {
        scanf("%d%d",&x,&y);
        e[++ne].to=y;e[ne].nxt=f1[x];f1[x]=ne;
        e[++ne].to=x;e[ne].nxt=f1[y];f1[y]=ne;
    }
    dfs1(1,0);
    for(i=1;i<=100000;++i)
        s+=t1[i];
    dfs2(1,0);
    for(i=1;i<=n;++i)
        printf("%lld
",ans[i]);
    return 0;
}