算法学习笔记(19): 树上启发式合并（DSU on tree）

树上启发式合并

DSU on tree，我也不知道DSU是啥意思

这是一种看似特别玄学的优化

可以把树上部分问题由 \(O(n^2)\) 优化到 \(O(n \log n)\)。

例如 CodeForces 600E。

又例如一道神奇的题：

适用情况

可以离线的部分树上问题。

需要子树上的所有信息，但是信息无法快速合并的情况。

或者说可以使用树上莫队的题目一般都可以使用启发式合并？（至少OI-Wiki是这么说的）

树上启发式合并并不是很常用

合并思路

首先定义一点概念：

• 重子节点：一个结点的子结点中度最大的结点（孩子个数最多的点）
• 轻子节点：非重子节点的节点
• 重边：重子结点（如果有）与其父结点相连的边
• 轻边：非重边
• 重链：相邻重边链接形成的链（这里好像用不到）

---

树上启发式合并的思路如下：

• 处理 x 为根的情况

• 递归处理 轻子节点 的答案，并舍弃其信息

• 处理 重子节点 的答案，并保留其信息

• 如果 x 为（其父亲的）重子节点，则加上所有 轻子节点 的信息。否则舍弃其信息

很明显，我们需要预处理出重子节点。同时可能需要用到 dfn 序来优化信息的加入

不妨我们以 vector 存图（非常方便）

void workSon(int x, int f) {
    siz[x] = 1, fa[x] = f;
    dfn[x] = ++cdfn, rdfn[cdfn] = x;
    for (int y : G[x]) {
        if (dfn[y]) continue;
        workSon(y, x);
        siz[x] += siz[y];
        if (siz[son[x]] < siz[y]) son[x] = y;
    }
    edfn[x] = cdfn;
}

最终 son[x] 就是 x 的重子节点，同时我们处理出了 dfn 序以及以 x 为子树的 dfn 序范围：[dfn[x], edfn[x]] 注意为闭区间

那么伪代码大致如下：

// remain 用于获知需不需要保留数据，默认不保留
int currentAnswer;
void work(int x, bool remain = false) {
    for (int y : G[x]) {
        if (y != fa[x] && y != son[x]) work(y);
    }

    if (son[x]) work(son[x], true);

    for (int y : G[x]) {
        if (y == fa[x] || y == son[x]) continue;
        addSubTreeInfoAndUpdateAnswer(y);
    }

    answerOf[x] = currentAnswer;
    if (!remain) {
        clearAnswer();
        for (int y : G[x]) {
            if (y != fa[x]) removeSubTreeInfo(y);
        }
    }
}

每个部分的作用在函数名里面应该很清晰了。这里不再赘述。

复杂度证明

首先，根节点到树上任意节点的轻边数不超过 \(\log n\) 条。

只有当祖先节点为轻子节点时其信息会被删除。也就是加入 \(O(\log n)\) 次，删除 \(O(\log n)\) 次，故而每一个点的复杂度为 \(O(\log n)\)，整体的复杂度为 \(O(n \log n)\)。

当然，考虑如果每一个节点加入信息或者删除信息的复杂度为 \(O(k)\)，则整体复杂度为 \(O(n k \log n)\)。

非常玄学……但是就是能够优化

例题分析

就以开始为例子的两道题为例吧。

CodeForces 600E

这道题也就是树上数颜色的问题。

题目大意是：

对于每一个节点，做出贡献的颜色需要满足出现的次数是最多的之一。一个颜色的贡献即是这个颜色的编号。

最终输出每一个节点被贡献的结果

样例输入
4
1 2 3 4
1 2
2 3
2 4
样例输出
10 9 3 4

最主要也是最重要的就是颜色的统计。

加入点（颜色）的核心代码如下：

int colorBut[N];
long long maxExi = 0, cnt = 0;
void add(int c) {
    if (++colorBut[c] == maxExi) {
        cnt += c;
    } else if (colorBut[c] > maxExi) {
        maxExi = colorBut[c], cnt = c;
    }
}

在合并部分也很清晰了

long long res[N];
void dsu(int x, int f, bool remain = false) {
    for (int y : G[x]) {
        if (y != f && y != son[x]) dsu(y, x);
    }

    if (son[x]) dsu(son[x], x, true);

    // 记得把根节点的信息也加进去！
    add(col[x]);
    for (int y : G[x]) {
        if (y == f || y == son[x]) continue;
        // 添加信息
        for (int i = dfn[y]; i <= edfn[y]; ++i) add(col[rdfn[i]]); 
    }

    res[x] = cnt;
    if (!remain) {
        maxExi = cnt = 0; // 重置答案
        for (int i = dfn[x]; i <= edfn[x]; ++i) // 删除影响信息
            colorBut[col[rdfn[i]]] = 0;
    }
}

不正常国家

这道题稍微复杂一点。

考虑每对点对其 LCA 的贡献。或者换个思路，考虑每一个根节点能够被那些点贡献。

不难发现，有两种情况：

• LCA 和其子节点之间的路径

• LCA 的两个子节点之间的路径。这里要保证两个子节点在不同的子树里面。

如果我们已经预处理出了树上异或前缀和 path。那么任意两个点对其 LCA 的贡献为 path[x] ^ path[y] ^ val[LCA]。我们不妨对于每一个 LCA，枚举所有 path[y] ^ val[LCA]，同时在已知的 path[x] 中匹配最大的异或对。

最大异或对可以看此题：AcWing 143.最大异或对

利用了01Trie树和二进制贪心。

此处不展开。

同时，由于我们需要保证 x 和 y 在 u 的不同子树中，所以我们先查询完一颗子树再加入这棵子树的信息。

核心代码如下：

// 树上启发式合并 
void work(int x, int f, bool remain = false) {
    // 首先搞定所有非重子节点
    for (int y : G[x]) {
        if (y == f || y == son[x]) continue;
        work(y, x);
    }

    // 搞定重子节点，并保留数据 
    if (son[x]) work(son[x], x, true);
    // path[fa[x]] 也就是 path[x] ^ val[x]
    int ans = max(val[x], trie.pairMax(path[fa[x]]));
    trie.insert(path[x]);

    // 加入其他节点，并搜索
    for (int y : G[x]) {
        if (y == f || y == son[x]) continue;
        for (int j = dfn[y]; j <= edfn[y]; ++j) {
            int pa = path[rdfn[j]] ^ val[x];
            ans = max(ans, trie.pairMax(pa));
        }

        for (int j = dfn[y]; j <= edfn[y]; ++j) {
            trie.insert(path[rdfn[j]]);
        }
    }

    res[x] = ans;
    if (!remain) trie.clear();
}

至于 01Trie 树代码如下：

const int LOG = 30; // 31位！下标为 [0, 30]

#define bit(x, i) ((x >> i) & 1)
class Trie01 {
private:
    int ch[N << 4][2];
    int usage;
public:
    Trie01() : usage(1) {
    }

    inline int newNode() {
        ++usage;
        ch[usage][0] = ch[usage][1] = 0;
        return usage; 
    }

    void insert(int x) {
        int p = 1;
        for (int k = LOG; k >= 0; --k) {
            int s = bit(x, k);
            if (!ch[p][s]) ch[p][s] = newNode();
            p = ch[p][s];
        }
    }

    // 这是通过树的形状贪心寻找最大异或对
    int pairMax(int x) {
        int p = 1;
        int r = 0;
        for (int k = LOG; k >= 0; --k) {
            int s = bit(x, k) ^ 1;
            if (ch[p][s]) r = (r << 1) | 1, p = ch[p][s];
            else if (ch[p][s ^ 1]) r <<= 1, p = ch[p][s ^ 1];
            else p = 0, r = x; // 避免空树的情况
        }
        return r;
    }

    void clear() {
        usage = 1;
        ch[1][0] = ch[1][1] = 0;
    }
} trie;

那么这道 水题 也就这么水过去了。

忘了说，其复杂度为 \(O(n \log n L)\)，其中 \(L\) 是位长，也就是代码中的 LOG = 30。所以复杂度也可以写为 \(O(n \log^2 n)\)

---

树上启发式合并的潜力不止于此，还望诸君发掘。