树上启发式合并（dsu on tree）学习笔记

闲话

树上启发式合并，又称 dsu on tree（虽然跟 dsu 并查集完全没关系），用于离线处理子树相关询问。

它是一种利用了重链剖分性质的暴力，时间复杂度为完全正确的 \(\mathcal{O}(n\log n+m)\)，个人认为跟莫队等都是非常优雅的暴力。

阅读本文并不需要重链剖分作为前置知识。

记号约定

本文中用 \(u\to v\) 表示 \(v\) 是 \(u\) 的儿子，\(u\leadsto v\) 表示 \(v\) 是 \(u\) 的后代。

例子：树上数颜色

这是树上启发式合并最基础的应用。

考虑暴力，每次询问搜一遍子树，统计答案，但这是 \(\mathcal{O}(mn)\) 的。如果预处理所有子树的信息，也是 \(\mathcal{O}(n^2+m)\) 的。

就没有更好的做法吗？

我们设 \({buc}_u\) 为 \(u\) 子树内的桶，\({buc}_{u,i}\) 也就是颜色 \(i\) 在 \(u\) 子树中出现的次数，也就是 \({buc}_{u,i}=\sum\limits_{u\leadsto v}[c_v=i]\)。

可以发现 \({buc}_{u,i}=\sum\limits_{u\to v}{buc}_{v,i}+[c_u=i]\)，但我们前面的暴力中每个节点都重复统计了 \(buc\)，是不是有些浪费呢？考虑对 \(buc\) 数组进行重复利用。

由于两棵没有包含关系的子树的 \(buc\) 不能重复利用，也不可能真的给每个节点都开一个 \(buc\)，因此对于每个节点，\(buc\) 数组只能从一个儿子处继承。那从哪里继承呢？用尾椎骨想一想就知道显然从重儿子处继承是最优的！这里我们把 \(u\) 的重儿子定义为，在 \(u\) 的所有儿子中，子树大小最大的。

于是算法流程就出来了：

1. 预处理每个点的重儿子，记点 \(u\) 的重儿子为 \({son}_u\)。
2. 从根递归下去，对于节点 \(u\)，先统计所有轻儿子（非重儿子）的答案，并从 \(buc\) 中擦去它们的贡献。此时的 \(buc\) 为空。
3. 然后统计重儿子的答案，但从 \(buc\) 中不擦去贡献。此时的 \({buc}_i=\sum\limits_{{son}_u\leadsto v}[c_v=i]\)。
4. 在 \(buc\) 中添加当前点 \(u\) 和所有轻儿子子树内的点的贡献。此时的 \({buc}_i=\sum\limits_{u\leadsto v}[c_v=i]\)。
5. 处理出当前点的答案。
6. 如果当前点是父亲的轻儿子，则需要擦去贡献，枚举 \(u\leadsto v\) 并把 \({buc}_{c_v}\) 减一即可。

那么复杂度是啥呢？

我们称一个节点与它的重儿子连接的边为“重边”，与它的轻儿子连接的边为“轻边”。

首先抛出一条引理：根节点到任意节点路径上的轻边不超过 \(\log n\) 条。

证明：设根节点到当前节点的轻边为 \(x\) 条，当前节点的子树大小为 \(y\)。由轻重儿子定义，显然轻儿子的子树大小不超过父亲的一半，则有 \(y < \dfrac{n}{2^x}\)，所以 \(n > 2^x\)，即 \(x < \log n\)。证毕。

然后考虑一个点会被计算多少次，显然只有搜到这个点，或者这个点位于它的某个祖先的轻儿子的子树内，这个点才会被计算。又因为上面的引理，显然每个点被计算次数为 \(\mathcal{O}(\log n)\)，于是复杂度为 \(\mathcal{O}(n\log n)\)。

回到这道树上数颜色，给出参考代码：

//By: Luogu@rui_er(122461)
#include <bits/stdc++.h>
#define rep(x,y,z) for(int x=y;x<=z;x++)
#define per(x,y,z) for(int x=y;x>=z;x--)
#define debug printf("Running %s on line %d...\n",__FUNCTION__,__LINE__)
#define fileIO(s) do{freopen(s".in","r",stdin);freopen(s".out","w",stdout);}while(false)
using namespace std;
typedef long long ll;
const int N = 1e5+5; 

int n, m, c[N], sz[N], son[N], buc[N], ans[N], now;
vector<int> e[N];
template<typename T> void chkmin(T& x, T y) {if(x > y) x = y;}
template<typename T> void chkmax(T& x, T y) {if(x < y) x = y;}
void dfs(int u, int f) { // 第一步：预处理重儿子
    sz[u] = 1;
    for(auto v : e[u]) {
        if(v == f) continue;
        dfs(v, u);
        sz[u] += sz[v];
        if(sz[v] > sz[son[u]]) son[u] = v;
    }
}
void add(int u, int f, int dt) {
    if(dt > 0 && !buc[c[u]]) ++now;
    buc[c[u]] += dt;
    if(dt < 0 && !buc[c[u]]) --now;
    for(auto v : e[u]) {
        if(v == f) continue;
        add(v, u, dt);
    }
}
void calc(int u, int f, int sv) {
    for(auto v : e[u]) { // 第二步：统计轻儿子答案并擦去贡献
        if(v == f || v == son[u]) continue;
        calc(v, u, 0);
    }
    if(son[u]) calc(son[u], u, 1); // 第三步：统计重儿子答案并保留贡献
    if(!buc[c[u]]) ++now;
    ++buc[c[u]];
    for(auto v : e[u]) { // 第四步：添加轻儿子贡献
        if(v == f || v == son[u]) continue;
        add(v, u, 1);
    }
    ans[u] = now; // 第五步：处理当前点答案
    if(!sv) add(u, f, -1); // 第六步：如果是轻儿子，擦去贡献
}

int main() {
    scanf("%d", &n);
    rep(i, 1, n-1) {
        int u, v;
        scanf("%d%d", &u, &v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    rep(i, 1, n) scanf("%d", &c[i]);
    dfs(1, 0);
    calc(1, 0, 0);
    for(scanf("%d", &m);m;m--) {
        int u;
        scanf("%d", &u);
        printf("%d\n", ans[u]);
    }
    return 0;
}

存几道题

CF600E Lomsat gelral、CF208E Blood Cousins、CF741D Arpa’s letter-marked tree and Mehrdad’s Dokhtar-kosh paths。

标签：子树,int,dsu,tree,儿子,启发式,buc,son,节点
来源： https://www.cnblogs.com/ruierqwq/p/dsu-on-tree.html