【洛谷2664】树上游戏（点分治）

点此看题面

大致题意： 给定一棵树，每个节点有一个颜色，定义\(s(i,j)\)为\(i\)到\(j\)路径上颜色数量，请你对于每一个\(i\)求出\(\sum_{j=1}^n s(i,j)\)。

点分治

这种题目比较显然是点分治吧... ...

大致思路

首先，按照点分治的基本套路，对于一棵子树内的路径，我们分两种情况讨论：

• 经过根节点的
• 不经过根节点的

呃，第二种情况就按照点分治继续处理下去就可以了，因此忽略不提。

而对于第一种情况的边，我们又可以分为两类：

• 以根节点为一个端点的
• 不以根节点为端点的

接下来就按此分类处理即可。

先是一波定义

首先，我们需要先来一波定义。

\(sum\)：所有颜色造成的总贡献。

\(col_i\)：节点\(i\)的颜色。

\(Val_i\)：颜色\(i\)对答案造成的贡献。

\(Size_i\)：以节点\(i\)为根的子树大小。

具体流程

• 首先是一遍\(dfs\)预处理。在这一遍与处理中，我们主要求出两个东西：\(Size\)和\(Val\)两个数组，当然，在求\(Val\)数组的同时，也要顺带求出\(sum\)。

• 然后，我们就可以对第一类路径进行处理，更新\(ans_{rt}\)。

  • 具体操作：\(ans_{rt}+=sum-Val_{col_{rt}}+Size_{rt}\)。

  • 这样做的理由是先将含有\(rt\)的颜色所造成的贡献从\(sum\)中删去（不然会造成重复计算），然后加上\(Size_{rt}\)（即含有\(rt\)的路径条数，\(rt\)的颜色在这些路径中会各对答案造成一点贡献）。

• 接下来便是对第二类路径的处理。我们考虑对于根节点的每一棵子树如何计算其贡献值。

  • 第一步自然是将这棵子树中的贡献清空，不然会造成重复计算。
  • 接下来，从该子节点开始，一直往下遍历，更新每一个节点的\(ans\)。
  • 首先，我们要将沿途经过的在该子树中出现的颜色的贡献从\(sum\)中减去，避免重复计算，并用一个变量\(color\_tot\)记录该子树中颜色种数。
  • 则对于当前节点\(x\)，我们可以将\(ans_x\)加上\(sum+color\_tot*Other\)。
  • 其中\(Other\)表示除该棵子树外其他子树以及根节点的总点数，因此将剩余贡献加上\(color\_tot*Other\)就是当前节点能造成的贡献了（和对第一类路径的操作同理）。

• 好吧，还有一个很重要的细节：注意清空数组。

  • 但是，如何清空数组也是有学问的。
  • 如果直接\(memset\)，显然复杂度成了\(O(n)\)，可能会\(TLE\)。
  • 所以，我们直接将该子树中出现的所有颜色的相关值全部清空即可，就是一个遍历子树的过程，时间复杂度为\(O(子树大小)\)，从而保证了时间复杂度。

代码

#include<bits/stdc++.h>
#define max(x,y) ((x)>(y)?(x):(y))
#define min(x,y) ((x)<(y)?(x):(y))
#define uint unsigned int
#define LL long long
#define ull unsigned long long
#define swap(x,y) (x^=y,y^=x,x^=y)
#define abs(x) ((x)<0?-(x):(x))
#define INF 1e9
#define Inc(x,y) ((x+=(y))>=MOD&&(x-=MOD))
#define ten(x) (((x)<<3)+((x)<<1))
#define N 100000
#define add(x,y) (e[++ee].nxt=lnk[x],e[lnk[x]=ee].to=y)
using namespace std;
LL n,ee=0,lnk[N+5],col[N+5];
struct edge
{
    LL to,nxt;
}e[(N<<1)+5];
class FIO
{
    private:
        #define Fsize 100000
        #define tc() (FinNow==FinEnd&&(FinEnd=(FinNow=Fin)+fread(Fin,1,Fsize,stdin),FinNow==FinEnd)?EOF:*FinNow++)
        #define pc(ch) (FoutSize<Fsize?Fout[FoutSize++]=ch:(fwrite(Fout,1,FoutSize,stdout),Fout[(FoutSize=0)++]=ch))
        LL f,FoutSize,OutputTop;char ch,Fin[Fsize],*FinNow,*FinEnd,Fout[Fsize],OutputStack[Fsize];
    public:
        FIO() {FinNow=FinEnd=Fin;}
        inline void read(LL &x) {x=0,f=1;while(!isdigit(ch=tc())) f=ch^'-'?1:-1;while(x=ten(x)+(ch&15),isdigit(ch=tc()));x*=f;}
        inline void read_char(char &x) {while(isspace(x=tc()));}
        inline void read_string(string &x) {x="";while(isspace(ch=tc()));while(x+=ch,!isspace(ch=tc())) if(!~ch) return;}
        inline void write(LL x) {if(!x) return (void)pc('0');if(x<0) pc('-'),x=-x;while(x) OutputStack[++OutputTop]=x%10+48,x/=10;while(OutputTop) pc(OutputStack[OutputTop]),--OutputTop;}
        inline void write_char(char x) {pc(x);}
        inline void write_string(string x) {register LL i,len=x.length();for(i=0;i<len;++i) pc(x[i]);}
        inline void end() {fwrite(Fout,1,FoutSize,stdout);}
}F;
class Class_DotSolver//点分治
{
    private:
        LL rt,sum,color_tot,used[N+5],Size[N+5],Max[N+5],cnt[N+5],Val[N+5];
        inline void GetRt(LL x,LL lst,LL tot)//找树的重心，确定根节点
        {
            register LL i;
            for(i=lnk[x],Size[x]=1,Max[x]=0;i;i=e[i].nxt)
                if(e[i].to^lst&&!used[e[i].to]) GetRt(e[i].to,x,tot),Size[x]+=Size[e[i].to],Max[x]=max(Max[x],Size[e[i].to]);
            if((Max[x]=max(Max[x],tot-Size[x]))<Max[rt]) rt=x;
        }
        inline void Clear(LL x,LL lst)//清空数组
        {
            for(register LL i=lnk[x];i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) Clear(e[i].to,x);//枚举子节点，遍历该子树
            cnt[col[x]]=Val[col[x]]=0;//清空与当前节点颜色相关的数据
        }
        inline void Init(LL x,LL lst)//初始化出子树大小和贡献
        {
            register LL i;
            for(i=lnk[x],Size[x]=1,++cnt[col[x]];i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) Init(e[i].to,x),Size[x]+=Size[e[i].to];//枚举子节点进行处理
            if(!--cnt[col[x]]) sum+=Size[x],Val[col[x]]+=Size[x];//如果该颜色是第一次出现，将当前颜色造成的贡献加上该子树大小
        }
        inline void Change(LL x,LL lst,LL flag)//删除（flag=0）or 恢复（flag=1）一棵子树中颜色造成的贡献
        {
            register LL i;
            for(i=lnk[x],++cnt[col[x]];i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) Change(e[i].to,x,flag);
            if(!--cnt[col[x]]) flag?(sum+=Size[x],Val[col[x]]+=Size[x]):(sum-=Size[x],Val[col[x]]-=Size[x]);//与上面的操作同理
        }
        inline void F5(LL x,LL lst,LL Other)//更新一棵子树中的ans
        {
            register LL i;
            if(!cnt[col[x]]++) sum-=Val[col[x]],++color_tot;//如果是第一次出现，将当前颜色的贡献减去，然后将该子树中颜色种数加1
            for(i=lnk[x],ans[x]+=sum+color_tot*Other;i;i=e[i].nxt) if(e[i].to^lst&&!used[e[i].to]) F5(e[i].to,x,Other);
            if(!--cnt[col[x]]) sum+=Val[col[x]],--color_tot;//撤销操作
        }
        inline void Solve(LL x)//处理以x为根的子树
        {
            register LL i,j;
            for(i=lnk[x],used[x]=1,Init(x,0),ans[x]+=sum-Val[col[x]]+Size[x];i;i=e[i].nxt)//初始化+更新根节点的ans
            {
                if(used[e[i].to]) continue;
                ++cnt[col[x]],sum-=Size[e[i].to],Val[col[x]]-=Size[e[i].to],Change(e[i].to,x,0),//先清空该子树造成的贡献
                --cnt[col[x]],F5(e[i].to,x,Size[x]-Size[e[i].to]),++cnt[col[x]],//更新该子树内ans
                Change(e[i].to,x,1),Val[col[x]]+=Size[e[i].to],sum+=Size[e[i].to],--cnt[col[x]];//恢复该子树造成的贡献
            }
            for(i=lnk[x],sum=color_tot=0,Clear(x,0);i;i=e[i].nxt)//清空
                if(!used[e[i].to]) GetRt(e[i].to,rt=0,Size[e[i].to]),Solve(rt);//继续处理子树信息
        }
    public:
        Class_DotSolver() {Max[0]=INF;}
        LL ans[N+5];
        inline void GetAns() {GetRt(1,rt=0,n),Solve(rt);}
}DotSolver;
int main()
{
    register LL i,x,y;
    for(F.read(n),i=1;i<=n;++i) F.read(col[i]);
    for(i=1;i<n;++i) F.read(x),F.read(y),add(x,y),add(y,x);
    for(DotSolver.GetAns(),i=1;i<=n;++i) F.write(DotSolver.ans[i]),F.write_char('\n');
    return F.end(),0;
}