[P2664] 树上游戏

Link：

P2664 传送门

Solution:

一道非常不错的计算贡献的题目

此类计算树上所有点对间结果的题目首先考虑点分治，同时一般都是对每种颜色计算贡献

那么对于每层点分治，需要计算其它点对分治中心的贡献和经过分支中心的路径对其它每个点的贡献

考虑颜色$k$，如果一个颜色为$k$的节点$v$为其到根的路径中的第一个$k$，则其对其它子树贡献为$size[v]$

令$col[k]$表示$sum size[v]$，$sum$表示$sum col[k]$

那么分治中心增加的答案就是$sum$，下面考虑每一颗子树中的点

首先要将子树中的点对$col[k]$和$sum$的贡献还原，在计算完该棵子树的贡献后再撤销该操作

接下来考虑每个点$v$，对其产生贡献的颜色$k$分两类：

1、不属于$v$到分治中心路径上的颜色，那么贡献就是$col[k]$

2、属于$v$到分治中心路径上的颜色，贡献显然为除了该子树的点的个数

因此在再次遍历时记录第二类点的个数同时不断用$sum$减去第一次出现颜色的$col[k]$即可

注意，每次解决完一个连通块要将数组还原

此时一定不能用$memset$，要记录经过的颜色并仅将这些还原，否则单层复杂度无法保证$O(n)$

Code:

#include <bits/stdc++.h>

using namespace std;
#define X first
#define Y second
typedef long long ll;
typedef pair<int,int> P;
const int MAXN=1e5+10;
struct edge{int nxt,to;}e[MAXN<<2];
ll cnt[MAXN],col[MAXN],res[MAXN],sum,num;
int n,x,y,vsum,vcur,rt,tot,top;
int head[MAXN],dat[MAXN],mxsub[MAXN],sz[MAXN],vis[MAXN],st[MAXN];

void add_edge(int from,int to)
{
    e[++tot].nxt=head[from];e[tot].to=to;head[from]=tot;
    e[++tot].nxt=head[to];e[tot].to=from;head[to]=tot;
}
void getrt(int x,int anc)
{
    sz[x]=1;mxsub[x]=0;
    for(int i=head[x];i;i=e[i].nxt)
        if(e[i].to!=anc&&!vis[e[i].to])
        {
            getrt(e[i].to,x);
            sz[x]+=sz[e[i].to];
            mxsub[x]=max(mxsub[x],sz[e[i].to]);
        }
    mxsub[x]=max(mxsub[x],vsum-sz[x]);
    if(mxsub[x]<mxsub[rt]) rt=x;
}
//计算总和 
void cal_sum(int x,int anc)
{
    sz[x]=1;cnt[dat[x]]++;
    for(int i=head[x];i;i=e[i].nxt)
        if(!vis[e[i].to]&&e[i].to!=anc)
            cal_sum(e[i].to,x),sz[x]+=sz[e[i].to];
    
    if(cnt[dat[x]]==1)
        col[dat[x]]+=sz[x],sum+=sz[x],st[++top]=dat[x];
    cnt[dat[x]]--;
}
void modify(int x,int anc,int f)
{
    cnt[dat[x]]++;
    for(int i=head[x];i;i=e[i].nxt)
        if(!vis[e[i].to]&&e[i].to!=anc) 
            modify(e[i].to,x,f);
    
    if(cnt[dat[x]]==1)
        sum+=sz[x]*f,col[dat[x]]+=sz[x]*f;
    cnt[dat[x]]--;
}
void cal_res(int x,int anc)//计算当前子树结果 
{
    cnt[dat[x]]++;
    if(cnt[dat[x]]==1)
        num++,sum-=col[dat[x]];
    res[x]+=sum+num*vcur;
    
    for(int i=head[x];i;i=e[i].nxt)
        if(!vis[e[i].to]&&e[i].to!=anc)
            cal_res(e[i].to,x);
    
    if(cnt[dat[x]]==1)
        num--,sum+=col[dat[x]];
    cnt[dat[x]]--;
}

void exclude(int x,int anc,int f)//将当前子树结果撤销/还原 
{
    sum+=sz[x]*f;col[dat[anc]]+=sz[x]*f;
    cnt[dat[anc]]++;modify(x,anc,f);cnt[dat[anc]]--;
}
void cal(int x)
{
    top=0;cal_sum(x,0);res[x]+=sum;
    for(int i=head[x];i;i=e[i].nxt)
    {
        if(vis[e[i].to]) continue;
        
        exclude(e[i].to,x,-1);
        vcur=sz[x]-sz[e[i].to];
        cal_res(e[i].to,x);
        exclude(e[i].to,x,1);
    }
    //一定不能用memset，仅对当前分治树的颜色清零 
    num=sum=0;
    for(int i=1;i<=top;i++) 
        cnt[st[i]]=col[st[i]]=0;
}
void solve(int x)
{
    cal(x);vis[x]=1;
    for(int i=head[x];i;i=e[i].nxt)
    {
        if(vis[e[i].to]) continue;
        vsum=mxsub[0]=sz[e[i].to];
        getrt(e[i].to,rt=0);solve(rt);
    }
}

int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n;i++) scanf("%d",&dat[i]);
    for(int i=1;i<n;i++)
        scanf("%d%d",&x,&y),add_edge(x,y);
    
    vsum=mxsub[0]=n;getrt(1,rt=0);
    solve(rt);
    for(int i=1;i<=n;i++)
        printf("%lld
",res[i]);
    return 0;
}