【学习笔记】Dsu On Tree

Dsu On Tree 是什么

针对树上问题的一种“优雅的暴力”，类似树上启发式合并。
基于轻重链剖分。

使用条件：

离线，无修改。
与子树有关，要遍历所有子树。

复杂度：\(O(nlogn)\)

解决问题示例

Problem

【 \(CF600E\) 】
一棵 \(n\) 个节点的树，每个结点有一个颜色。
问，以每个节点为根的子树中，出现次数最多的颜色是哪种（若多种，输出种类编号之和）。

朴素的暴力

以每个点为根，遍历每棵子树统计答案。
复杂度 \(O(n^2)\)

优雅的暴力

在朴素的暴力中，许多重复的信息没有重复利用。
比如，当遍历完一棵子树后，我们将它对答案的贡献清除了。但其实在遍历其父节点的子树时，这棵子树的信息是有用的。

在递归 \(dfs\) 中，可不可以将某些子树的答案保留，再将其它子树的答案暴力并上去，得到父节点的答案？
当然 \(ok\)
而那棵被保留的子树，贪心地想应是节点数最多的子树，即“重子”。

于是基本思路就出来了——
先做一遍轻重链剖分。然后在递归时先递归轻子，之后将答案清除；再递归重子，将答案保留；随后将轻点的答案暴力加上，统计本节点的答案。
这就是 \(dsu\) \(on\) \(tree\) ！撒花！

复杂度证明

轻重链剖分有个性质：每个点到根的路径上最多只有 \(logn\) 条轻边。
证明：每走过一条轻边，子树大小一定小于之前的一半。把 \(n\) 减半，最多减 \(logn\) 次。

本算法中，只有轻边所连的子树会被暴力统计。
考虑每个点被暴力统计多少次，即其在多少个轻边所连子树中，其实就是它到根节点的轻边个数，即 \(logn\)。
故，总复杂度 \(nlogn\)

代码

#include<cstdio>
#include<iostream>
#include<algorithm>

using namespace std;

int read(){
    int x=0;
    char ch=getchar();
    while(!isdigit(ch)) ch=getchar();
    while(isdigit(ch)) x=x*10+ch-'0',ch=getchar();
    return x;
}

const int N = 100005;
typedef long long ll;

struct node{
    int v;
    node *nxt;
}pool[N*2],*h[N];
int cnt;
void addedge(int u,int v){
    node *p=&pool[++cnt],*q=&pool[++cnt];
    p->v=v;p->nxt=h[u];h[u]=p;
    q->v=u;q->nxt=h[v];h[v]=q;
}

int sz[N],son[N],fa[N];
void dfs1(int u){
    int v,Mson=0;
    sz[u]=1;
    for(node *p=h[u];p;p=p->nxt)
        if(!fa[v=p->v]){
            fa[v]=u;
            dfs1(v);
            sz[u]+=sz[v];
            if(sz[v]>Mson) Mson=sz[v],son[u]=v;
        }
}

int n,c[N],t[N],mx;
ll ans[N],num;

void Plus(int x,int y) { t[x]+=y; if(t[x]==mx) num+=x; if(t[x]>mx) mx=t[x],num=x; }
void add(int u,int x){
    int v;
    Plus(c[u],x);
    for(node *p=h[u];p;p=p->nxt)
        if(fa[v=p->v]==u) add(v,x);
}
void dfs(int u,int ty){
    int v;
    for(node *p=h[u];p;p=p->nxt)
        if(fa[v=p->v]==u && v!=son[u]) dfs(v,0);
    if(son[u]) dfs(son[u],1);
    for(node *p=h[u];p;p=p->nxt) 
        if(fa[v=p->v]==u && v!=son[u]) add(v,1);
    Plus(c[u],1);
    ans[u]=num;
    if(ty==0) add(u,-1),mx=0,num=0;
}

int main()
{
    n=read();
    for(int i=1;i<=n;i++) c[i]=read();
    for(int i=1;i<n;i++) addedge(read(),read());
    
    fa[1]=-1; dfs1(1);
    dfs(1,1);
    
    for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
    
    return 0;
}

算法（伪）代码

//轻重链剖分
void dfs1(int u){
    int v,Mson=0;
    sz[u]=1;
    for(node *p=h[u];p;p=p->nxt)
        if(!fa[v=p->v]){
            fa[v]=u;
            dfs1(v);
            sz[u]+=sz[v];
            if(sz[v]>Mson) Mson=sz[v],son[u]=v;
        }
}
//dsu on tree
void Plus(int u,int x);
void add(int u,int x){ //暴力枚举轻子子树求贡献
        Plus(u,x);
        for(node *p=h[u];p;p=p->nxt) add(p->v,x);
}
void cal();
void work(int u,int ty){
        int v;
        for(node *p=h[u];p;p=p->nxt)
            if((v=p->v)!=son[u]) work(v,0);//递归处理轻子,处理后清理
        if(son[u]) work(son[u],1);//递归处理重子，之后不清理
        
        for(node *p=h[u];p;p=p->nxt)
            if((v=p->v)!=son[u]) add(v,1); //暴力加上轻子的答案
        Plus(u,1);//加上本节点的答案
        
        cal(); //统计

        if(ty==0) {  //若本节点为轻点，则需清除贡献
            add(u,-1); 
            与答案有关变量=0；//能全清为零，是由于轻点在重点之前递归，所有贡献均为该轻点内的贡献，而无需保留的重点的贡献。
        }
}

其他例题及应用

（详见 \(cf\) 博客）
\(cf570D\) 也是个模板题
还可以解决一些点分治的问题。