题解 P2664 【树上游戏】

· · 题解

考虑点分治，对于当前分治中心，统计出它自己出发到分治块内的所有路径对自己答案的贡献，和经过它的路径对当前分治块内点的贡献。自己出发到分治块内的所有路径对自己答案的贡献很好求，现在考虑怎么求经过它的路径对当前分治块内点的贡献。

我们对于当前分治中心的每一个子树分别考虑，令cnt[i]为从分治中心出发的进入其他子树的所有路径中，包含颜色i的路径条数，size为除了该子树外当前分治块内所有的点的个数。那么，我们dfs这棵子树计算贡献，假设当前dfs到x，首先给sum_x加上\sum cnt[i]，即所有在其他子树中出现的颜色的贡献总和，然后计算x到分治中心的路径上颜色的贡献。

对于一个出现在分治中心到x的路径上的颜色c，它对x的贡献为size-cnt[c]，因为c已经在一些路径上出现，它现在能产生的额外贡献为它原来没有出现的路径条数。所以我们给sum_x还要加上size-cnt[c]。同时c也会对x的子树内所有点产生贡献，所以这个贡献要像标记一样往下传递，然后标记一下c的贡献已经被计算过，往下dfs时就不用再次计算了。

所以只需要对于当前分治中心求出cnt数组和每棵子树的size，进入一棵子树时减去子树自己内部对cnt产生的贡献。同时为了防止复杂度退化，我们不能对于所有颜色求cnt，要先统计一下当前分治块内有哪些颜色出现了，这样枚举块内所有颜色的复杂度才是O(分治块大小)。

复杂度O(n\log n)。

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;

#define For(i,_beg,_end) for(int i=(_beg),i##end=(_end);i<=i##end;++i)
#define Rep(i,_beg,_end) for(int i=(_beg),i##end=(_end);i>=i##end;--i)

template<typename T>T Max(const T &x,const T &y){return x<y?y:x;}
template<typename T>T Min(const T &x,const T &y){return x<y?x:y;}
template<typename T>int chkmax(T &x,const T &y){return x<y?(x=y,1):0;}
template<typename T>int chkmin(T &x,const T &y){return x>y?(x=y,1):0;}
template<typename T>void read(T &x){
    T f=1;char ch=getchar();
    for(;ch<'0'||ch>'9';ch=getchar())if(ch=='-')f=-1;
    for(x=0;ch>='0'&&ch<='9';ch=getchar())x=x*10+ch-'0';
    x*=f;
}

typedef long long LL;
const int maxn=100010,inf=0x7fffffff;
struct edge{
    int to,nxt;
}e[maxn<<1];
int n,num,head[maxn],c[maxn];
int f[maxn],size[maxn],Size,root;
int has[maxn],col[maxn],top,cl[maxn],tp;
LL sum[maxn],res[maxn],cnt[maxn],ct[maxn],cct[maxn],path,tot;
bool vis[maxn],exist[maxn];

void addedge(int,int);
void get_sum(int,int);//计算子树大小
void get_root(int,int);//计算重心
void Dfs(int,int,LL*);//统计哪些颜色出现过，同时计算cnt数组
void Modify(int,int,LL);//更新子树的答案
void Calc(int);//对于分治中心统计答案
void Solve(int);//点分治

int main(){
    read(n);
    For(i,1,n) read(c[i]);
    For(i,1,n-1){
        int u,v;
        read(u);read(v);
        addedge(u,v);
    }
    Solve(1);
    For(i,1,n) printf("%lld\n",sum[i]);
    return 0;
}

void Calc(int x){
    get_sum(x,Size=0);

    tot=top=0;
    Dfs(x,0,cnt);
    For(i,1,top) exist[col[i]]=false;//要及时还原标记数组
    For(i,1,tp=top){
        tot+=cnt[cl[i]=col[i]];
        cct[col[i]]=cnt[col[i]];
    }

    sum[x]+=tot;
    LL temp=tot;

    for(int i=head[x];i;i=e[i].nxt) if(!vis[e[i].to]){

        has[c[x]]=1;top=0;
        Dfs(e[i].to,x,ct);
        has[c[x]]=0;
        For(j,1,top) exist[col[j]]=false;

        cnt[c[x]]-=size[e[i].to];
        tot-=size[e[i].to];
        For(j,1,top){
            cnt[col[j]]-=ct[col[j]];
            tot-=ct[col[j]];
        }

        path=size[x]-size[e[i].to];
        Modify(e[i].to,x,0);

        cnt[c[x]]+=size[e[i].to];
        tot=temp;
        For(j,1,top){
            cnt[col[j]]=cct[col[j]];
            ct[col[j]]=0;
        }
    }

    For(i,1,tp) cnt[cl[i]]=0;//还原数组
    vis[x]=true;
}
void Modify(int x,int fa,LL lst){//lst为之前下传下来的贡献和
    LL tag=lst;
    if(++has[c[x]]==1) tag+=path-cnt[c[x]];
    sum[x]+=tot+tag;
    for(int i=head[x];i;i=e[i].nxt){
        if(e[i].to==fa||vis[e[i].to]) continue;
        Modify(e[i].to,x,tag);
    }
    has[c[x]]--;
}
void Dfs(int x,int fa,LL *cnt){
    if(!exist[c[x]]) col[++top]=c[x],exist[c[x]]=true;
    if(++has[c[x]]==1) cnt[c[x]]+=size[x];
    for(int i=head[x];i;i=e[i].nxt){
        if(e[i].to==fa||vis[e[i].to]) continue;
        Dfs(e[i].to,x,cnt);
    }
    has[c[x]]--;
}
void Solve(int x){
    Calc(x);
    for(int i=head[x];i;i=e[i].nxt) if(!vis[e[i].to]){
        f[Size=root=0]=inf;
        get_sum(e[i].to,x);
        get_root(e[i].to,x);
        Solve(root);
    }
}
void get_sum(int x,int fa){
    Size++;size[x]=1;
    for(int i=head[x];i;i=e[i].nxt){
        if(e[i].to==fa||vis[e[i].to]) continue;
        get_sum(e[i].to,x);
        size[x]+=size[e[i].to];
    }
}
void get_root(int x,int fa){
    f[x]=0;
    for(int i=head[x];i;i=e[i].nxt){
        if(e[i].to==fa||vis[e[i].to]) continue;
        get_root(e[i].to,x);
        chkmax(f[x],size[e[i].to]);
    }
    chkmax(f[x],Size-size[x]);
    if(f[x]<f[root]) root=x;
}
void addedge(int u,int v){
    e[++num].to=v;e[num].nxt=head[u];head[u]=num;
    e[++num].to=u;e[num].nxt=head[v];head[v]=num;
}