树上启发式合并/dsu on tree

树上启发式合并/dsu on tree

前置芝士

启发式合并和树链剖分的部分知识。（不会的去这里搜）

因为要在一颗树上进行启发式合并，所以要找最优的方法，即优雅的暴力（雾

它可以让 \(O(n^2)\) 变为 \(O(n\log n)\) (证明 你就想想启发式合并就完了)

概念

树上启发式合并（dsu on tree）对于某些树上离线问题可以速度大于等于大部分算法且更易于理解和实现的算法。

因为没什么需要讲的，所以直接看题

例题1

U41492 树上数颜色

题意

给一棵根为1的树，每次询问子树颜色种类数

思路

​ 直接暴力预处理的时间复杂度为 O(n^2) ，即对每一个子节点进行一次遍历，每次遍历的复杂度显然与 n 同阶，有 n 个节点，故复杂度为 O(n^2)

​ 可以发现，每个节点的答案是其子树的叠加，考虑利用这个性质处理问题

​ 我们可以先预处理出每个节点子树的size和它的重儿子，重儿子同树链剖分一样，是拥有节点最多子树的儿子，这个过程显然可以O(n)完成

​ 我们用check[i]表示颜色 i 有没有出现过，ans[i]表示他的颜色个数

​ 遍历一个节点，我们按以下的步骤进行遍历：

• 先遍历其非重儿子，获取它的ans，但不保留

• 遍历后它的check遍历它的重儿子，保留它的check

• 再次遍历其非重儿子及其父亲，用重儿子的check对遍历到的节点进行计算，获取整棵子树的ans

  为什么不合并第一步和第三步呢？因为check数组不能重复使用，否则空间会太大，需要在 O(n) 的空间内完成。

  代码

  #include<bits/stdc++.h>
  using namespace std;
  #define ll int
  #define r register 
  #define A 1001010
  ll head[A],nxt[A],ver[A],size[A],col[A],cnt[A],ans[A],son[A];
  ll tot=0,num,sum,nowson,n,m,xx,yy;
  inline void add(ll x,ll y){
      nxt[++tot]=head[x],head[x]=tot,ver[tot]=y;
  }
  inline ll read(){
      ll f=1,x=0;char c=getchar();
      while(!isdigit(c)){
          if(c=='-') f=-1;
          c=getchar();
      }
      while(isdigit(c))
          x=(x<<1)+(x<<3)+(c^48),c=getchar();
      return f*x;
  }
  void dfs(ll x,ll fa){
      size[x]=1;
      for(ll i=head[x];i;i=nxt[i]){
          ll y=ver[i];
          if(y==fa) continue;
          dfs(y,x);
          size[x]+=size[y];
          if(size[son[x]]<size[y])
              son[x]=y;
      }
  }
  void cal(ll x,ll fa,ll val){
      if(!cnt[col[x]]) ++sum;
      cnt[col[x]]+=val;
      for(ll i=head[x];i;i=nxt[i]){
          ll y=ver[i];
          if(y==fa||y==nowson) continue;
          cal(y,x,val); 
      }
  }
  void dsu(ll x,ll fa,bool op){
      for(ll i=head[x];i;i=nxt[i]){
          ll y=ver[i];
          if(y==fa||y==son[x])
              continue;
          dsu(y,x,0);
          //从轻儿子出发
      }
      if(son[x])
          dsu(son[x],x,1),nowson=son[x];
      cal(x,fa,1);nowson=0;
      ans[x]=sum;
      if(!op){
          cal(x,fa,-1);
          sum=0;
      }
  }
  int main(){
      n=read();
      for(ll i=1;i<=n-1;i++){
          xx=read(),yy=read();
          add(xx,yy),add(yy,xx);
      }
      for(ll i=1;i<=n;i++)
          col[i]=read();
      dfs(1,0);
      dsu(1,0,1);
      m=read();
      for(ll i=1;i<=m;i++){
          xx=read();
          printf("%d\n",ans[xx]);
      }
  }
  

例题2

CF600E Lomsat gelral

思路

这道题我们可以遍历整棵树，并用一个数组ap（appear）记录每种颜色出现几次

先用一遍dfs算出每个点是否为重儿子

再dfs统计答案，每次碰到重儿子就跳过，递归完清空ap数组等东东

最后dfs重儿子，不清空

再对当前节点进行另一种dfs，暴力统计ap，不做重儿子

代码

#include<bits/stdc++.h>
#define ll long long
#define re register
using namespace std;
const int N=2e5+10;
int n;
int c[N];//color
int v[N],nex[N],first[N],tot=1;
inline void add(int x,int y){
    v[++tot]=y;
    nex[tot]=first[x];
    first[x]=tot;
}
inline int read(){
    int x=0;char ch=getchar();
    while(!isdigit(ch))ch=getchar();
    while(isdigit(ch))x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
    return x;
}
ll ans[N],ap[N],mx,sum;//十年OI一场空，不开 long long 见祖宗 
//ap表示每种颜色出现几次 mx表示出现最多的次数 sum表示颜色编号和 
int sz[N];//子树大小 
bool gson[N];//表示一个点是否为重儿子 
void getg(int x,int f){//get 子树大小 以及 重儿子 
    sz[x]=1;
    int mx=0,p=0;
    for(re int i=first[x];i;i=nex[i]){
        int y=v[i];
        if(y==f)continue;
        getg(y,x);sz[x]+=sz[y];
        if(sz[y]>mx){
            mx=sz[y];
            p=y;
        }
    }if(p)gson[p]=1;
}
void DFS(int x,int f,int p){//暴力遍历子树 p为重儿子 之后需init清空 
    //统计答案 
    ap[c[x]]++;
    if(ap[c[x]]>mx){
        mx=ap[c[x]];
        sum=c[x];
    }else if(ap[c[x]]==mx)sum+=c[x];
    for(re int i=first[x];i;i=nex[i]){
        int y=v[i];
        if(y==f || y==p)continue;//不要把重儿子也一起遍历了！ 
        DFS(y,x,p);
    }
}
inline void init(int x,int f){//暴力遍历后清空 
    ap[c[x]]--;
    for(re int i=first[x];i;i=nex[i]){
        int y=v[i];
        if(y==f)continue;
        init(y,x);
    }
}
void dfs(int x,int f){//启发式合并关键函数！ 
    int p=0;//重儿子标记 
    for(re int i=first[x];i;i=nex[i]){
        int y=v[i];
        if(y==f)continue;
        if(!gson[y]){//不是重儿子的暴力做 
            dfs(y,x);
            init(y,x);
            sum=mx=0;
        }
        else p=y;
    }if(p)dfs(p,x);//重儿子单独特判
    DFS(x,f,p);
    ans[x]=sum;
}
int main(){
    n=read();
    for(re int i=1;i<=n;i++)c[i]=read();
    for(re int i=1;i<n;i++){
        int x=read(),y=read();
        add(x,y),add(y,x);
    }getg(1,0);
    dfs(1,0);
    for(re int i=1;i<=n;i++)
        printf("%lld ",ans[i]);
}

（自己写的程序WA了，现在只能拿一个题解的代码）