P2664 树上游戏 （点分治）

题目链接

题面：

题解：
点分治。
考虑当前点分治子树的根节点为p。
那么我们需要处理两部分贡献。
①、所有节点对点p的贡献（包括p节点自己）
②、经过点p的路径（x，y）对x点的贡献。

第一类贡献：
如果一个节点 x 的颜色 color [ x ] 在它到根节点 p 的路径上是第一次出现，那么以 x 为根的子树上的每个节点，都可以对节点 p 产生贡献，那么一共产生了 size [ x ] 的贡献。
为方便我们最后计算和统计答案，我们把所有这样的贡献加起来记为sum。显然该当前分治子树上所有节点对p的贡献和就是sum。即 ans[p]+=sum。

void dfs_sum(int x,int fa)
{
    si[x]=1;
    cnt[vi[x]]++;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]||y==fa) continue;
        dfs_sum(y,x);
        si[x]+=si[y];
    }
    if(cnt[vi[x]]==1)
    {
        sum+=si[x];
        c[vi[x]]+=si[x];
    }
    cnt[vi[x]]--;
}

第二类贡献。
假设当前考虑点经过p的路径对点x的贡献
我们将第二类贡献分为两部分

第一部分：
路径p–x出现的颜色对x点的贡献（不包括p点）。这一部分的贡献为 (size[ p ] - size[ sonp ] ) * res
其中 sonp 为 p 的儿子且为 x 的祖先，res为p–x出现的颜色种类数（不包括p点）。

第二部分：
除了 sonp 以外子树上新出现的颜色对 x 的贡献，我们求出了sum，所有节点对p的贡献，那么我们需要减去sonp子树中对sum的贡献，然后再减去p–x出现的颜色对sum的贡献（不包括p点），此时的sum就是sonp以外子树新出现的颜色对x的贡献

void dfs_ans(int x,int fa)
{
    cnt[vi[x]]++;
    if(cnt[vi[x]]==1)
    {
        sum-=c[vi[x]];
        res++;
    }
    ans[x]+=sum+res*six_xiy;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(y==fa||ha[y]) continue;
        dfs_ans(y,x);
    }
    if(cnt[vi[x]]==1)
    {
        sum+=c[vi[x]];
        res--;
    }
    cnt[vi[x]]--;
}

代码：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<string>
#include<cmath>
#include<algorithm>
#include<cstdlib>
#define ll long long
#define llu unsigned ll
#define pr make_pair
#define pb push_back
using namespace std;
const int maxn=100100;
const int inf=0x3f3f3f3f;
int head[maxn],ver[maxn<<1],nt[maxn<<1];
int max_part[maxn],si[maxn],vi[maxn],c[maxn],cnt[maxn];
bool ha[maxn];
int tot=1,root,max_size,n;
ll ans[maxn],sum,res,six_xiy;

void add(int x,int y)
{
    ver[++tot]=y,nt[tot]=head[x],head[x]=tot;
}

void dfs_size(int x,int fa)
{
    si[x]=1;max_part[x]=0;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]||y==fa) continue;
        dfs_size(y,x);
        si[x]+=si[y];
        max_part[x]=max(max_part[x],si[y]);
    }
}

void dfs_root(int now_root,int x,int fa)
{
    max_part[x]=max(max_part[x],si[now_root]-si[x]);
    if(max_size>max_part[x])
    {
        max_size=max_part[x];
        root=x;
    }
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]||y==fa) continue;
        dfs_root(now_root,y,x);
    }
}

void dfs_sum(int x,int fa)
{
    si[x]=1;
    cnt[vi[x]]++;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]||y==fa) continue;
        dfs_sum(y,x);
        si[x]+=si[y];
    }
    if(cnt[vi[x]]==1)
    {
        sum+=si[x];
        c[vi[x]]+=si[x];
    }
    cnt[vi[x]]--;
}

void change(int x,int fa,int val)
{
    cnt[vi[x]]++;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(y==fa||ha[y]) continue;
        change(y,x,val);
    }
    if(cnt[vi[x]]==1)
    {
        sum+=si[x]*val;
        c[vi[x]]+=si[x]*val;
    }
    cnt[vi[x]]--;
}

void dfs_ans(int x,int fa)
{
    cnt[vi[x]]++;
    if(cnt[vi[x]]==1)
    {
        sum-=c[vi[x]];
        res++;
    }
    ans[x]+=sum+res*six_xiy;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(y==fa||ha[y]) continue;
        dfs_ans(y,x);
    }
    if(cnt[vi[x]]==1)
    {
        sum+=c[vi[x]];
        res--;
    }
    cnt[vi[x]]--;
}

void dfs_clear(int x,int fa)
{
    cnt[vi[x]]=0;
    c[vi[x]]=0;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(y==fa||ha[y]) continue;
        dfs_clear(y,x);
    }
}

void do_it(int x)
{
    sum=0,res=0;
    dfs_sum(x,0);
    ans[x]+=sum;
    for(int i=head[x];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]) continue;
        cnt[vi[x]]++;
        sum-=si[y];
        c[vi[x]]-=si[y];
        change(y,x,-1);
        cnt[vi[x]]--;

        six_xiy=si[x]-si[y];
        dfs_ans(y,x);

        cnt[vi[x]]++;
        sum+=si[y];
        c[vi[x]]+=si[y];
        change(y,x,1);
        cnt[vi[x]]--;
    }
    dfs_clear(x,0);
}

void dfs(int x)
{
    max_size=inf;
    dfs_size(x,0);
    dfs_root(x,x,0);
    do_it(root);
    ha[root]=true;
    for(int i=head[root];i;i=nt[i])
    {
        int y=ver[i];
        if(ha[y]) continue;
        dfs(y);
    }
}


int main(void)
{
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
        scanf("%d",&vi[i]);
    int x,y;
    for(int i=1;i<n;i++)
        scanf("%d%d",&x,&y),add(x,y),add(y,x);
    dfs(1);
    for(int i=1;i<=n;i++)
        printf("%lld\n",ans[i]);
    return 0;
}