DSU on tree（树上启发式合并）

适用范围/算法特点

• 树上启发式合并是一种处理树上不带修改的离线子树问题的算法
• 时间复杂度为

前置知识

• 无

问题引入

先来看cf上的一道不算模版的题：1709-E XOR Tree - Codeforces
给定有 () 个节点的树，每个节点有一个数字 () 。一条简单路径的权重定义为路径上所有点的异或值。如果没有简单路径的权重为，则这棵树称为好的。
你可以进行任意次操作，每次操作修改一个节点的值为任意值。求为了使这棵树变成“好的”，最少需要的操作次数。
输入
第一行一个数，表示节点个数
第二行n个整数, , ..., ()---节点的数字。
接下来行，每行两个数、，表示一条连接点、的边。

input
6
3 2 1 3 2 1
4 5
3 4
1 4
2 1
6 1
output
2

Solution

• 对于路径权值的问题，让我们取1为根节点。首先可以求出从到的路径权重。假设一条路径上的两个端点为和，而他们的最近公共祖先为，那么这条路径的权值等于就等价于 ^ ^ == 。整理得 ^ == 。从而处理成子树的问题。
• 对于修改操作，我们可以将这个节点的值改为一个很大的值，从而使这个数的某一位二进制只有它为1（因为修改的值没有上限），这样经过这个点的路径权值都不为了。

我们发现如果一棵树去暴力处理它的所有子树，时间复杂度为。DSU on tree优秀的地方就在于此。我们每个节点开一个set，将一个节点的所有信息都insert到这个节点。我们发现这个时候交换两个节点的set是不会影响的，因此我们可以让该节点的set成为较大的那个set。可以证明，这样的时间复杂度为。而对于一旦发现需要修改，直接clear即可。

Code

#include<bits/stdc++.h>
#define endl '\n'
using namespace std;
using ll = long long;

void solve() {
    int n;
    cin >> n;
    vector<int> c(n + 1);
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
    }
    vector<vector<int>> E(n + 1, vector<int>());
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        E[u].push_back(v);
        E[v].push_back(u);
    }
    vector<int> d(n + 1);
    
    // dfs0处理出d[u]
    function<void(int u, int fa)> dfs0 = [&](int u, int fa) {
        d[u] = d[fa] ^ c[u];
        for (auto v : E[u]) {
            if (v == fa) continue;
            dfs0(v, u);
        }
    };
    
    set<int> s[n + 1];
    int ans = 0;
    
    // 这里u是正在遍历的节点, v是其子节点, fa是其父节点
    function<void(int, int)> dfs = [&](int u, int fa) {
        s[u].insert(d[u]);
        bool f = 0; // f打标记, 若需要清空就置1
        for (auto v : E[u]) {
            if (v == fa) continue;
            dfs(v, u);

            // 最关键的一步, 可以发现后面的两次都是遍历s[v]的, 因此我们把小的set给v
            if (s[u].size() < s[v].size()) {
                swap(s[u], s[v]);
            }

            for (auto x : s[v]) {
                if (s[u].count(x ^ c[u])) f = 1;
            }
            for (auto x : s[v]) {
                s[u].insert(x);
            }
        }
        
        // 注意细节, 这里是全部遍历再清空set
        if (f) {
            s[u].clear();
            ans++;
        }
    };
    dfs0(1, 0);
    dfs(1, -1);
    cout << ans << endl;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    int tt = 1;
    // cin >> tt;
    while (tt--) solve();
    return 0;
}

ok成功AC

正式介绍

其实上面的例子不算是DSU on tree算法的例子，但它却处处存在着该算法的影子。可以发现，这是一种“优美的暴力”。我们暴力处理每个节点的子节点，但是时间复杂度却来到了。这是由于我们都只处理了size较小的set。相对于上题的代码，DSU on tree给出了三点扩展:
1. 我们不必给每一个节点开一个状态（即set）；（1）
2. 每个点的暴搜过程可以直接用dfs遍历一遍；（2）
3. 交换set大小的操作改为先遍历重儿子，再遍历轻儿子。（3）
（重儿子是指该节点size最大的儿子，其他都为轻儿子）

下面先给出一道真正的模板题，以上的特点会一一解释。

600-E Lomsat gelral - Codeforces
给定有 () 个节点的有根树，根节点为。每个节点有一个颜色。对于每个节点，求出他的子树中数量最多的颜色。如果有多种数量最多的颜色，求出他们的总和。
输入
第一行一个数，表示节点个数
第二行n个整数, , ..., ()---节点的数字。
接下来行，每行两个数、，表示一条连接点、的边。
输出
个数，第个数表示以为根的子树中数量最多的颜色的总和。

input
4
1 2 3 4
1 2
2 3
2 4
output
10 9 3 4

下面的思路参考了【AgOHの算法胡扯】树上启发式合并（dsu on tree）_哔哩哔哩_bilibili

算法流程
对于节点：
* 处理好所有轻儿子及其子树的贡献，算出来后删除轻儿子的贡献也就是说先遍历轻儿子；
* 处理好重儿子及其子树的贡献，算完后不删除贡献对应；
* 暴力统计所有轻儿子及其子树的贡献，就是此时节点的答案对应。

对于，我们可以开一个总的数组，这样在遍历完节点的时候这个数组就是它的答案。这就是为什么我们在遍历完轻儿子的时候要删除掉它的贡献，因为它的贡献会影响到其他子树。
对于，重儿子是后遍历的，因此它可以保存下来。而它的size又是最大的，这样时间复杂度最优。
其实思路和上题XOR Tree非常相似。

具体细节可以看code

Code

#include<bits/stdc++.h>
#define endl '\n'
using namespace std;
using ll = long long;

void solve() {
    int n;
    cin >> n;
    vector<int> col(n + 1);
    vector<ll> ans(n + 1);
    vector<vector<int>> E(n + 1, vector<int>());
    for (int i = 1; i <= n; i++) {
        cin >> col[i];
    }
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        E[u].push_back(v);
        E[v].push_back(u);
    }

    vector<int> sz(n + 1), son(n + 1);

    // dfs0处理出所有重儿子和轻儿子
    function<void(int, int)> dfs0 = [&](int u, int fa) {
        sz[u] = 1;
        for (auto v : E[u]) {
            if (v == fa) continue;
            dfs0(v, u);
            sz[u] += sz[v];
            if (sz[v] > sz[son[u]]) {
                son[u] = v;
            }
        }
    };

    // 这里cnt记录了每个颜色的个数, 将数组开在外面
    vector<int> cnt(n + 1);
    ll sum = 0, maxc = 0;
    // flag的作用后面会提到
    int flag = 0;

    // 下面的update函数是暴搜的过程, val分别取1和-1, 代表加贡献/删贡献
    function<void(int, int, int)> update = [&](int u, int fa, int val) {
        // 以下八行计算贡献, 不同的题目需求不同
        cnt[col[u]] += val;
        if (cnt[col[u]] > maxc) {
            maxc = cnt[col[u]];
            sum = col[u];
        }
        else if (cnt[col[u]] == maxc) {
            sum += col[u];
        }

        // 这里为什么不写v == son[u]呢？因为只有一个重儿子son[u]不需要遍历, 那就是最开始那个u
        // 所以开一个flag, flag就是最开始的flag, 不理解的可以看下面函数的调用时机
        for (auto v : E[u]) {
            if (v == fa || v == flag) continue;
            update(v, u, val);
        }
    };

    // 第三个参数keep表示是否需要清空
    function<void(int, int, bool)> dfs = [&](int u, int fa, bool keep) {
        for (auto v : E[u]) {
            if (v == fa || v == son[u]) continue;
            // 轻儿子需要清空, keep为false
            dfs(v, u, false);
        }
        if (son[u]) {
            // 重儿子不需要清空, keep为true
            dfs(son[u], u, true);
            flag = son[u];
        }

        // 可以看到在这里只需要统计轻儿子就可以, 只有当前u的重儿子不需要遍历, 而轻儿子可能有与他自己的重儿子, 还是需要遍历的, 因此需要开一个flag
        update(u, fa, 1);
        flag = 0;

        // 此时所有儿子都已经遍历完一遍, 当前的cnt数组就是该节点的子树的所有状态
        ans[u] = sum;
        if (!keep) {
            update(u, fa, -1);
            maxc = sum = 0;
        }
    };
    dfs0(1, -1);
    dfs(1, -1, 0);
    for (int i = 1; i <= n; i++) {
        cout << ans[i] << " ";
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    int tt = 1;
    // cin >> tt;
    while (tt--) solve();
    return 0;
}

DSU on tree题单：洛谷 | DSU on tree