【洛谷P2664】 树上游戏 点分治

code:

#include <bits/stdc++.h>

#define N 200009

#define ll long long

#define setIO(s) freopen(s".in","r",stdin)

using namespace std;

ll Sum[N];

int n,edges,root,sn;

int val[N],hd[N],to[N<<1],nex[N<<1],size[N],mx[N],vis[N],A[N];

inline void add(int u,int v)

{

    nex[++edges]=hd[u],hd[u]=edges,to[edges]=v;

}

void getroot(int u,int ff)

{

    size[u]=1,mx[u]=0;

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(v==ff||vis[v])    continue;

        getroot(v,u);

        size[u]+=size[v];

        mx[u]=max(mx[u],size[v]);

    }

    mx[u]=max(mx[u],sn-size[u]);

    if(mx[u]<mx[root])   root=u;

}

int ou;

ll tmp,tot,bu[N];

map<int,ll>cn[N];

map<int,ll>::iterator it;

int dep[N],cnt[N],siz[N];

void getnode(int top,int u,int ff,int cur)

{

    if(!cnt[val[u]])

        ++cur;

    ++cnt[val[u]];

    Sum[top]+=(ll)cur;

    siz[u]=1;

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(v==ff||vis[v])     continue;

        getnode(top,v,u,cur);

        siz[u]+=siz[v];

    }

    --cnt[val[u]];

}

void get_col(int top,int u,int ff)

{

    if(!cnt[val[u]])

        cn[top][val[u]]+=(ll)siz[u];

    ++cnt[val[u]];

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(v==ff||vis[v])   continue;

        get_col(top,v,u);

    }

    --cnt[val[u]];

}

void calc_v(int u,int ff)

{

    ll tt=bu[val[u]];

    tmp=tmp-bu[val[u]]+ou;

    bu[val[u]]=ou;

    Sum[u]+=tmp;

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(vis[v]||v==ff)    continue;

        calc_v(v,u);

    }

    tmp=tmp-bu[val[u]]+tt;

    bu[val[u]]=tt;

}

void clr(int u,int ff)

{

    cn[u].clear();

    bu[val[u]]=0;

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(v==ff||vis[v])   continue;

        clr(v,u);

    }

}

void calc(int u)

{

    tot=0;

    getnode(u,u,0,0);

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(vis[v])    continue;

        // memset(cnt,0,sizeof(cnt));

        get_col(v,v,u);

        for(it=cn[v].begin();it!=cn[v].end();it++)

        {

            tot+=it->second;

            bu[it->first]+=it->second;

        }

    }

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(vis[v])    continue;

        tmp=tot;

        ou=siz[u]-siz[v];

        for(it=cn[v].begin();it!=cn[v].end();it++)

        {

            bu[it->first]-=it->second;

            tmp-=it->second;

        }

        ll tt=bu[val[u]];

        tmp=tmp-bu[val[u]]+ou;

        bu[val[u]]=ou;

        calc_v(v,u);

        bu[val[u]]=tt;

        for(it=cn[v].begin();it!=cn[v].end();it++) bu[it->first]+=it->second;

    }

    clr(u,0);

}

void dfs(int u)

{

    calc(u);

    vis[u]=1;

    for(int i=hd[u];i;i=nex[i])

    {

        int v=to[i];

        if(vis[v])    continue;

        root=0,sn=size[v],getroot(v,u),dfs(root);

    }

}

int main()

{

    // setIO("input");

    int i,j;

    scanf("%d",&n);

    for(i=1;i<=n;++i)     scanf("%d",&val[i]), A[i]=val[i];

    sort(A+1,A+1+n);

    for(i=1;i<=n;++i)     val[i]=lower_bound(A+1,A+1+n,val[i])-A;

    for(i=1;i<n;++i)

    {

        int u,v;

        scanf("%d%d",&u,&v),add(u,v),add(v,u);

    }

    sn=mx[0]=n,root=0,getroot(1,0),dfs(root);

    for(i=1;i<=n;++i)     printf("%lld\n",Sum[i]);

    return 0;

}