dsu on tree 学习笔记

发布时间 2023-06-28 17:43:48作者: 霜木_Atomic

dsu on tree 学习笔记

引入

dsu 是并查集的缩写,然鹅本算法和并查集没啥关系。当然,dsu on tree 也有中文名字:树上启发式合并。也就是说,这个算法是用于处理一些树上信息的合并的。

dsu on tree 和莫队一样,都是优雅的暴力。优雅是因为思想很优雅,暴力是因为所有修改都是暴力做的——没错,遍历整棵子树的那种暴力。

思想

当然,虽然是暴力,直接搞肯定会 T。既然是要合并信息,我们肯定是要能少删除就少删除,也就是说,保留信息最多的,让其他信息往上合并。

dsu on tree 利用了重链剖分的一个性质来实现这一点。我们每次都先处理轻儿子,再把轻儿子的信息删除,然后再处理重儿子,保留重儿子信息的同时再去暴力把轻儿子的信息求出来,以达到求出这一棵子树中所有节点答案的目的。

看着很暴力,但是实际上,这样做的复杂度只有 \(O(n \log_n)\)!为什么呢?我们知道,在重链剖分中,每一个轻儿子最多跳 \(\log_n\) 条轻边就会跳到重链上。我们可以反过来想:轻儿子的信息什么时候不用被清空呢?也就是当轻儿子所在的重儿子(就是重链上的一部分)被处理到时。那么,我们有这样一个结论:每一个节点的信息最多被清空 \(\log_n\) 次,最终的复杂度就是 \(O(n \log_n)\)

例题

CF600E Lomsat gelral

模板题,具体见代码注释。

#include<bits/stdc++.h>
using namespace std;
const int N = 1e5+100;

inline int read(){
    int x = 0; char ch = getchar();
    while(ch<'0' || ch>'9') ch = getchar();
    while(ch>='0'&&ch<='9') x = x*10+ch-48, ch = getchar();
    return x;
}
int head[N], tot;
struct node{
    int nxt, to;
}edge[N<<1];
void add(int u, int v){
    edge[++tot].nxt = head[u];
    edge[tot].to = v;
    head[u] = tot;
}

int col[N];
int n;

long long vis[N];
long long sum_ans, now_ans;
long long ans[N];

int siz[N], son[N];
void dfs1(int u, int fath){//和树链剖分的dfs1基本一样,求出重儿子
    siz[u] = 1;
    for(int i = head[u]; i; i = edge[i].nxt){
        int v = edge[i].to;
        if(v == fath) continue;
        dfs1(v, u);
        siz[u]+=siz[v];
        if(siz[son[u]]<siz[v]) son[u] = v;
    }
}

void update(int u, int fath){//增加信息,更新答案
    vis[col[u]]++;
    if(vis[col[u]]>sum_ans) now_ans =col[u], sum_ans = vis[col[u]];
    else if(vis[col[u]]==sum_ans) now_ans+=col[u];//可能有不止一种优势颜色
    for(int i = head[u]; i; i = edge[i].nxt){
        int v = edge[i].to;
        if(v == fath){
            continue;
        }
        update(v, u);
    }
}

void del(int u, int fath){//删除信息,清空轻儿子的贡献
    vis[col[u]]--;
    for(int i = head[u]; i; i = edge[i].nxt){
        int v = edge[i].to;
        if(v == fath){
            continue;
        }
        del(v, u);
    }
}
void dfs(int u, int fath, bool op){
    for(int i = head[u]; i; i = edge[i].nxt){
        int v = edge[i].to;
        if((fath ^ v)&&(son[u] ^ v)){
            dfs(v, u, 1);
        }//先搞轻儿子
    }
    if(son[u]) dfs(son[u], u, 0);//最后搞重儿子
    vis[col[u]]++;//注意当前节点的贡献不要忘了加上
    if(vis[col[u]]>sum_ans) now_ans =col[u], sum_ans = vis[col[u]];
    else if(vis[col[u]]==sum_ans) now_ans+=col[u];

    for(int i = head[u]; i; i = edge[i].nxt){//此时重儿子已经处理好且信息未删除,现在把轻儿子的贡献合并上来
        int v = edge[i].to;
        if((v ^ fath) && (v ^ son[u])){
            update(v, u);
        }
    }
    ans[u] = now_ans;//记录答案
    if(op) now_ans = sum_ans = 0, del(u, fath);//如果是轻儿子就清空当前子树的贡献
}
int main(){
    n = read();
    for(int i = 1; i<=n; ++i) col[i] = read();
    for(int i = 1; i<n; ++i){
        int u = read(), v = read();
        add(u, v);
        add(v, u);
    }
    dfs1(1, 0);
    dfs(1, 0, 0);
    for(int i = 1; i<=n; ++i){
        printf("%lld", ans[i]);
        if(i<n) putchar(' ');
        else putchar('\n');
    }
    return 0;
}