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

点此看题面

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

---

点分治

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

$Link$

点分治 详见博客 初学点分治

---

大致思路

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

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

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

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

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

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

---

先是一波定义

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

$sum$：所有颜色造成的总贡献。

$co{l}_i$：节点$i$的颜色。

$Va{l}_i$：颜色$i$对答案造成的贡献。

$Siz{e}_i$：以节点$i$为根的子树大小。

---

具体流程

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

• 然后，我们就可以对第一类路径进行处理，更新$an{s}_{rt}$。

  • 具体操作：$an{s}_{rt}+=sum-Va{l}_{co{l}_{rt}}+Siz{e}_{rt}$。

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

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

  • 第一步自然是将这棵子树中的贡献清空，不然会造成重复计算。
  • 接下来，从该子节点开始，一直往下遍历，更新每一个节点的$ans$。
  • 首先，我们要将沿途经过的在该子树中出现的颜色的贡献从$sum$中减去，避免重复计算，并用一个变量$color\_tot$记录该子树中颜色种数。
  • 则对于当前节点$x$，我们可以将$an{s}_x$加上$sum+color\_tot\ast Other$。
  • 其中$Other$表示除该棵子树外其他子树以及根节点的总点数，因此将剩余贡献加上$color\_tot\ast 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;
}