树上游戏

这里是一个辣鸡的\(\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;
}

转载于:https://www.cnblogs.com/AH2002/p/10420036.html