樹上游戲

這裏是一個辣雞的\(\Theta(n \log^2n)\)作法

首先咱們套個點分治

而後問題就變成了統計跨過根的路徑的貢獻

先跑一遍以當前點爲根的子樹大小和DFS序

同時將根路徑上首次出現該顏色的點存起來

而後根據顏色排序

對於不一樣顏色統計結果顯然互不影響

那麼對於同種顏色咱們容斥一下就能夠了

區間加法能夠用DFS序上差分來代替

#include"cstdio"
#include"cstring"
#include"iostream"
#include"algorithm"
using namespace std;

const int MAXN=1e5+5;

int n,np,rt,sum,tp;
int h[MAXN];
int siz[MAXN],mx[MAXN],sz[MAXN],id[MAXN],rid[MAXN];
int clr[MAXN],bnk[MAXN],sm[MAXN];
long long v[MAXN];
bool vis[MAXN];
long long ans[MAXN];
struct rpg{
    int li,nx;
}a[MAXN<<1];
struct pt{
    int x,c,r;
}stk[MAXN];

int read()
{
    int x=0;char ch=getchar();
    while(ch<'0'||'9'<ch) ch=getchar();
    while('0'<=ch&&ch<='9') x=(x<<3)+(x<<1)+(ch^48),ch=getchar();
    return x;
}

void add(int ls,int nx)
{
    a[++np]=(rpg){h[ls],nx};h[ls]=np;
    a[++np]=(rpg){h[nx],ls};h[nx]=np;
}

void gtrt(int x,int fa)
{
    siz[x]=1;mx[x]=0;
    for(int i=h[x];i;i=a[i].li){
        if(a[i].nx==fa||vis[a[i].nx]) continue;
        gtrt(a[i].nx,x);
        siz[x]+=siz[a[i].nx];
        mx[x]=max(mx[x],siz[a[i].nx]);
    }mx[x]=max(mx[x],sum-siz[x]);
    if(mx[rt]>mx[x]) rt=x;
    return;
}

void gtsz(int x,int fa)
{
    sz[x]=1;id[x]=++id[0];rid[id[0]]=x;
    for(int i=h[x];i;i=a[i].li){
        if(a[i].nx==fa||vis[a[i].nx]) continue;
        gtsz(a[i].nx,x);
        sz[x]+=sz[a[i].nx];
    }return;
}

void dfs(int x,int fa,int rot)
{
    if(!bnk[clr[x]]) stk[++tp]=(pt){x,clr[x],rot};
    ++bnk[clr[x]];
    for(int i=h[x];i;i=a[i].li){
        if(a[i].nx==fa||vis[a[i].nx]) continue;
        dfs(a[i].nx,x,rot);
    }--bnk[clr[x]];
    return;
}

bool cmp(pt a,pt b){return a.c<b.c;}

void calc(int x)
{
    id[0]=0;
    gtsz(x,x);bnk[clr[x]]=1;
    ans[x]+=sz[x];//由根致使的起點爲根的路徑對根的貢獻
    for(int i=h[x];i;i=a[i].li){
        if(vis[a[i].nx]) continue;
        dfs(a[i].nx,a[i].nx,a[i].nx);
        v[id[a[i].nx]]+=sz[x]-sz[a[i].nx];
        v[id[a[i].nx]+sz[a[i].nx]]-=sz[x]-sz[a[i].nx];
            //由根致使的起點在該子樹內跨過根的路徑對起點的貢獻
    }sort(stk+1,stk+tp+1,cmp);
    for(int l=1,r=0;l<=tp;l=r+1){
        int sum=0;
        while(r+1<=tp&&stk[r+1].c==stk[l].c) ++r,sm[stk[r].r]+=sz[stk[r].x],sum+=sz[stk[r].x];
        for(int i=l;i<=r;++i){
            v[id[stk[i].x]]-=sum-sm[stk[i].r];
            v[id[stk[i].x]+sz[stk[i].x]]+=sum-sm[stk[i].r];
            //減去跨過根的相同顏色致使多算的貢獻
            v[1]+=sz[stk[i].x];
            v[id[stk[i].r]]-=sz[stk[i].x];
            v[id[stk[i].r]+sz[stk[i].r]]+=sz[stk[i].x];
            //統計終點在該子樹內路徑對起點的貢獻
            v[id[stk[i].x]]+=sz[x]-sz[stk[i].r];
            v[id[stk[i].x]+sz[stk[i].x]]-=sz[x]-sz[stk[i].r];
                //起點在該子樹內對起點的貢獻
        }for(int i=l;i<=r;++i) sm[stk[i].r]=0;
    }
    for(int i=1;i<=id[0];++i) v[i]+=v[i-1],ans[rid[i]]+=v[i];
    for(int i=1;i<=id[0]+1;++i) v[i]=0;
    bnk[clr[x]]=tp=0;
    return;
}

void DFZ(int x)
{
    vis[x]=1;calc(x);
    for(int i=h[x];i;i=a[i].li){
        if(vis[a[i].nx]) continue;
        sum=siz[a[i].nx];rt=0;
        gtrt(a[i].nx,a[i].nx);
        DFZ(rt);
    }return;
}

int main()
{
    n=read();
    for(int i=1;i<=n;++i) clr[i]=read();
    for(int i=1;i<n;++i){
        int x=read(),y=read();
        add(x,y);
    }sum=n;mx[rt]=n;
    gtrt(1,1);
    DFZ(rt);
    for(int i=1;i<=n;++i) printf("%lld\n",ans[i]);
    return 0;
}