我的算法书 - 树上启发式合并

一句话算法

树上启发式合并保留重儿子的统计结果，只把轻儿子的贡献重新加入，从而避免反复清空大块信息。

问题模型

给定一棵以 1 为根的树，每个点有一个颜色。对每个节点 u，统计它的子树中出现次数最多的颜色。

如果有多个颜色出现次数相同，模板输出这些颜色编号之和。

这个模型来自经典的 DSU on tree 入门题：

子树查询很多；
每个查询都需要统计一大块信息；
信息不能简单用一个值合并，例如颜色频次表。

核心直觉

如果对每个节点都重新扫一遍子树，总复杂度可能是 $O(n^2)$ 。

树上启发式合并的思路是：

先找出每个节点的重儿子，也就是子树最大的儿子。
处理轻儿子时，算完就清空它的统计。
处理重儿子时，保留它的统计。
最后把所有轻儿子的贡献加入重儿子的统计表中。

这样，每个点只会在跨过轻边时被重新加入。树上一条根到叶路径最多有 $O(\log n)$ 条轻边，所以总复杂度是 $O(n\log n)$ 。

算法步骤

第一次 DFS：求重儿子

计算 subtree_size[u]。
找到 u 的最大子树儿子，记为 heavy_child[u]。

第二次 DFS：启发式合并

对节点 u：

先递归处理所有轻儿子，参数 keep=false，处理完清空统计。
再递归处理重儿子，参数 keep=true，保留统计。
标记重儿子为 big_child，避免重复加入。
把 u 和所有轻儿子子树的贡献加入当前统计表。
当前统计表就是 u 子树的统计结果。
如果 keep=false，离开 u 时清空整棵子树贡献。

算法证明

核心不变量：执行完 dfs_solve(u) 中的加入步骤后，统计表中恰好包含 u 子树内所有点的贡献。

轻儿子先被独立处理，答案已经记录；因为它们不会被父节点直接复用，所以清空不影响正确性。

重儿子最后处理并保留统计。由于重儿子子树是 u 子树的一部分，保留它可以避免重复加入最大的一块。

随后把 u 本身和所有轻儿子子树加入统计表，此时统计表包含：

重儿子子树；
所有轻儿子子树；
节点 u 本身。

这正好是 u 的整棵子树。

每个点在被清空后，只有当它属于某个祖先的轻儿子子树时才会被重新加入。沿根到该点的路径，每经过一条轻边，所在子树大小至少减半，因此轻边数量最多 $O(\log n)$ 。所以总加入次数为 $O(n\log n)$ 。

复杂度分析

第一次 DFS： $O(n)$ 。
第二次 DFS： $O(n\log n)$ 。
空间复杂度： $O(n+C)$ ，其中 $C$ 是颜色值域大小。

若题目颜色值很大，需要先离散化颜色。

代码实现

模板输入格式：

n
c1 c2 ... cn
n-1 条边

输出每个节点子树中出现次数最多的颜色编号之和。

cpp

        1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

        #include <bits/stdc++.h>
using namespace std;

struct DsuOnTree {
    int n = 0;
    vector<vector<int>> graph;
    vector<int> color;
    vector<int> subtree_size;
    vector<int> heavy_child;
    vector<int> answer;
    vector<int> color_count;
    int max_count = 0;
    int sum_color = 0;
    int big_child = 0;

    DsuOnTree(int n, int max_color) : n(n) {
        graph.assign(n + 1, {});
        color.assign(n + 1, 0);
        subtree_size.assign(n + 1, 0);
        heavy_child.assign(n + 1, 0);
        answer.assign(n + 1, 0);
        color_count.assign(max_color + 1, 0);
    }

    void add_edge(int u, int v) {
        graph[u].push_back(v);
        graph[v].push_back(u);
    }

    void dfs_size(int u, int parent) {
        subtree_size[u] = 1;
        for (int v : graph[u]) {
            if (v == parent) continue;
            dfs_size(v, u);
            subtree_size[u] += subtree_size[v];
            if (subtree_size[v] > subtree_size[heavy_child[u]]) {
                heavy_child[u] = v;
            }
        }
    }

    void add_color(int c, int delta) {
        color_count[c] += delta;
        if (delta > 0) {
            if (color_count[c] > max_count) {
                max_count = color_count[c];
                sum_color = c;
            } else if (color_count[c] == max_count) {
                sum_color += c;
            }
        }
    }

    void add_subtree(int u, int parent, int delta) {
        add_color(color[u], delta);
        for (int v : graph[u]) {
            if (v == parent || v == big_child) continue;
            add_subtree(v, u, delta);
        }
    }

    void reset_state() {
        fill(color_count.begin(), color_count.end(), 0);
        max_count = 0;
        sum_color = 0;
    }

    void dfs_solve(int u, int parent, bool keep) {
        for (int v : graph[u]) {
            if (v == parent || v == heavy_child[u]) continue;
            dfs_solve(v, u, false);
        }

        if (heavy_child[u]) {
            dfs_solve(heavy_child[u], u, true);
            big_child = heavy_child[u];
        }

        add_subtree(u, parent, 1);
        big_child = 0;
        answer[u] = sum_color;

        if (!keep) {
            add_subtree(u, parent, -1);
            reset_state();
        }
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;

    vector<int> input_color(n + 1);
    int max_color = 0;
    for (int i = 1; i <= n; i++) {
        cin >> input_color[i];
        max_color = max(max_color, input_color[i]);
    }

    DsuOnTree solver(n, max_color);
    for (int i = 1; i <= n; i++) {
        solver.color[i] = input_color[i];
    }

    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        solver.add_edge(u, v);
    }

    solver.dfs_size(1, 0);
    solver.dfs_solve(1, 0, true);

    for (int i = 1; i <= n; i++) {
        if (i > 1) cout << ' ';
        cout << solver.answer[i];
    }
    cout << '\n';

    return 0;
}

测试用例

输入：

输出：

3 2 1 2 3

解释：

节点 1 的子树中颜色 1 和 2 都出现两次，答案为 1+2=3。
节点 2 的子树中颜色 2 出现两次，答案为 2。

应用分类详解

树上启发式合并适合处理“每个子树都要维护一张较大统计表”的问题。

一、子树颜色统计

典型模式： 每个节点询问自己的子树中颜色频次、众数、不同颜色数。

识别信号： 出现“对每个节点的子树统计”“颜色出现次数”。

核心建模： DFS 序或子树遍历负责加入点，统计表维护颜色信息。

应用场景	经典题目	核心思路
子树众数	cf-600E	保留重儿子统计，加入轻儿子
子树颜色种数	树上统计题	维护颜色计数和非零颜色数

二、子树集合信息

典型模式： 每个子树对应一个集合，集合合并代价较高。

识别信号： 出现“每个节点的子树集合”“合并 map/set”。

核心建模： 小集合并入大集合，减少元素搬运次数。

三、离线树上查询

典型模式： 查询固定在子树上，不涉及在线修改。

识别信号： 所有查询提前给出，树结构不变。

核心建模： 按 DFS 顺序处理节点，统计表只在遍历过程中维护。

经典例题

1. cf-600E

Lomsat gelral。树上启发式合并经典题，统计每个子树中出现次数最多的颜色编号和。

2. 子树颜色种数

把统计表改成颜色出现次数，并维护当前出现次数大于 0 的颜色数量。

3. 子树最大频次问题

维护颜色频次、最大频次以及达到最大频次的颜色集合信息。