[蓝桥杯 2023 省 A] 颜色平衡树

形式化题目

给定一棵以 $1$ 号点为根、共 $n$ 个结点的树（父亲编号 $F_i < i$ ），每个结点 $i$ 有一个颜色 $C_i$ 。

对每个结点 $u$ ，记它的子树结点集合为 $S_u$ ，子树大小为 $S = |S_u|$ ；记子树中出现过的颜色种数为 $k$ ，颜色 $c$ 在 $S_u$ 中出现的次数为 $\mathrm{cnt}_u(c)$ 。

若 $\mathrm{cnt}_u(c_1) = \mathrm{cnt}_u(c_2) = \cdots = \mathrm{cnt}_u(c_k)$ ，其中 $c_1,\dots,c_k$ 是 $S_u$ 中出现过的全部颜色，则称子树 $u$ 是颜色平衡的。

求颜色平衡的子树的个数。

需要注意：参与比较的只有出现过的颜色。子树里没出现的颜色不算“出现 0 次”去和别的颜色比较，所以“子树里只有一种颜色”一定平衡。

暴力解法

思路

题意的逐字翻译：对每个结点 $u$ ，重新把子树 $u$ 完整走一遍，用桶统计每种颜色出现几次，再看这些出现过的颜色次数是否全部相同。

具体做法：

用 bucket[c] 记录颜色 $c$ 在当前子树中出现几次。
从 $u$ 出发递归遍历整棵子树，把每个点的颜色扔进桶，同时把“第一次出现的颜色”记到 seen[] 里。
遍历 seen[] 中的所有颜色，检查它们的 bucket 值是否全部相等。
用 seen[] 把桶清干净，恢复给下一个结点使用。

这个写法完全不需要任何额外观察，是最容易确认正确的基线，专门用来对拍和理解题意。

代码

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/**
 * Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
 * rbook: -> https://rbook.roj.ac.cn  https://rbook2.roj.ac.cn
 * rainboy的学习导航网站: https://idx.roj.ac.cn
 * create_at: 2026-10-02 15:07
 * update_at: 2026-10-03 11:07
 */
// brute.cpp：小数据暴力解，用来帮助理解题意并辅助对拍。
// 做法：对每个结点 u，重新完整遍历一次子树 u，统计每种颜色出现几次，
// 再检查所有出现过的颜色点数是否相同。复杂度约为 O(n^2)，只适合小数据。
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int MAXN = 200005;

int n;
int col[MAXN];              // col[i]：结点 i 的颜色
int fa[MAXN];               // fa[i]：结点 i 的父亲
vector<int> son_list[MAXN]; // 孩子列表

int bucket[MAXN];            // bucket[c]：本轮统计中颜色 c 的结点个数
int seen[MAXN];              // seen[1..seen_cnt]：本轮出现过的颜色，用来检查和清空桶
int seen_cnt;

ll ans;

// 把 u 及其所有后代的颜色扔进桶，返回这棵子树的结点个数。
int collect(int u) {
    if (bucket[col[u]] == 0) {
        seen_cnt++;
        seen[seen_cnt] = col[u];
    }
    bucket[col[u]]++;

    int total = 1;
    for (int i = 0; i < (int)son_list[u].size(); i++) {
        total += collect(son_list[u][i]);
    }
    return total;
}

// 对每个结点暴力统计它自己的子树。
void dfs_solve(int u) {
    seen_cnt = 0;
    collect(u);

    // 所有出现过的颜色，结点个数必须完全一样。
    bool ok = true;
    int first = bucket[seen[1]];
    for (int i = 2; i <= seen_cnt; i++) {
        if (bucket[seen[i]] != first) {
            ok = false;
            break;
        }
    }
    if (ok) {
        ans++;
    }

    // 还原桶，供后面的结点使用。
    for (int i = 1; i <= seen_cnt; i++) {
        bucket[seen[i]] = 0;
    }

    for (int i = 0; i < (int)son_list[u].size(); i++) {
        dfs_solve(son_list[u][i]);
    }
}

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

    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> col[i] >> fa[i];
    }
    for (int i = 2; i <= n; i++) {
        son_list[fa[i]].push_back(i);
    }

    dfs_solve(1);
    cout << ans << "\n";

    return 0;
}

复杂度

设树的结点数为 $n$ 。对每个结点 $u$ ，第 2 步要遍历整棵子树（ $O(\mathrm{siz}[u])$ ），第 3、4 步要遍历出现过的颜色（不超过 $\mathrm{siz}[u]$ ），所以单个结点是 $O(\mathrm{siz}[u])$ 。

总复杂度为

\sum_{u=1}^{n} O(\mathrm{siz}[u]) = O\!\left(\sum_{u=1}^{n} \mathrm{siz}[u]\right)

而 $\sum_u \mathrm{siz}[u] = O(n^2)$ （一条长度为 $n$ 的链上， $\mathrm{siz}[u] = n-u+1$ ，求和约 $n^2/2$ ）。

于是暴力是 $O(n^2)$ ：只能通过 $30\%$ 的数据（ $n \leqslant 200$ ）；对 $n \leqslant 5000$ 尚能勉强通过；对 $n = 2\times10^5$ 完全不可行。在链状数据上的实测（brute.cpp，GCC 16 -O2）能看清这个增长趋势：

$n$	10000	20000	40000
用时	0.41 s	1.61 s	6.31 s

每翻倍一次耗时约四倍，到 $n = 2\times10^5$ 就要约 150 秒量级。

瓶颈

瓶颈非常明确：每个结点都从零重新统计自己的子树，同一个点的颜色会被它的每一个祖先重复数一遍，重复次数等于它的深度。

想要变快，就必须让父亲复用儿子的统计结果，把“每个点被数几次”从“深度”降下来。

正解

思路

第一步：把判定改成 $O(1)$

暴力里最贵的部分之一是第 3 步“遍历出现过的颜色看是否全相等”。但其实不需要真的逐个比较。

设子树 $u$ 的大小为 $S$ ，出现过的颜色种数为 $k$ 。如果子树平衡，设公共次数为 $t$ ，那么这 $k$ 种颜色一共覆盖了 $S$ 个点：

S = k \cdot t \quad\Longrightarrow\quad t = \frac{S}{k}

反过来，如果 $k \mid S$ ，并且恰好有 $k$ 种颜色的出现次数等于 $S/k$ ，那么这 $k$ 种颜色必然就是全部出现过的颜色（因为出现过的颜色一共只有 $k$ 种），于是每种都恰好 $S/k$ 次，子树平衡。

所以我们只需要维护一个额外的数组：

occ[t]：当前统计范围内，恰好出现 $t$ 次的颜色种数。

再配合 kind（出现过的颜色种数），判定就变成了两行：

cpp

1if (siz[u] % kind == 0 && occ[siz[u] / kind] == kind) ans++;

occ 和 kind 都可以在“把一个点加入/移出统计”时 $O(1)$ 维护：

把一个颜色 $c$ 的点加入时，它的次数从 $x$ 变成 $x+1$ ：occ[x]--，occ[x+1]++；若 $x = 0$ ，说明这是一个新颜色，kind++。
移出时对称：occ[x]--；若 $x - 1 > 0$ 则 occ[x-1]++，否则这个颜色消失了，kind--。

occ[0] 没有意义，代码里用 cnt[c] > 0 来判断“这个颜色是否出现过”，就不必维护 occ[0]。

第二步：让父亲复用儿子的统计

现在我们有了 $O(1)$ 的判定，但还需要“统计范围”正好等于子树。最直接的想法是让 $u$ 继承所有儿子的统计，但这行不通：统计状态全局只有一份，算完儿子 $v_1$ 再算 $v_2$ 时，必须先清掉 $v_1$ 的贡献，否则判 $v_2$ 会被污染。清掉的开销是 $O(\mathrm{siz}[v_1])$ 。

如果每个儿子都清掉，那每个点还是会被反复计入，复杂度退化成 $O(n^2)$ 。

DSU on Tree（树上启发式合并） 的关键动作是：只留一个“值得留”的儿子。

每个结点选出子树最大的儿子，称为重儿子 heavy[u]，其余儿子称为轻儿子。
处理顺序：先把所有轻儿子各自算完、算完立刻清空；再算重儿子，算完保留；最后把轻儿子子树和 $u$ 自己补进来。此时统计范围恰好是子树 $u$ ，直接判定，然后按需清空。

这样父亲手上有重儿子现成的统计，只需要额外加入“轻儿子们的总大小 + 1”个点，而不是重新加入一整棵子树。

第三步：为什么这样是 $O(n\log n)$

需要回答一个问题：一个点到底会被加入多少次？

引理：若 $v$ 是 $u$ 的轻儿子，则 $\mathrm{siz}[v] \leqslant \mathrm{siz}[u] / 2$ 。

理由是：重儿子是 $u$ 的所有儿子中子树最大的那个。如果 $\mathrm{siz}[v] > \mathrm{siz}[u]/2$ ，那么 $u$ 的其它子树加起来只有 $\mathrm{siz}[u] - 1 - \mathrm{siz}[v] < \mathrm{siz}[u]/2 < \mathrm{siz}[v]$ ， $v$ 就会是重儿子，矛盾。

于是从任意结点往上走到根，每经过一条轻边，所在子树大小至少翻倍。子树大小不超过 $n$ ，所以一条路径上的轻边最多 $\lfloor \log_2 n \rfloor$ 条。

而一个点 $x$ 被“整体重新加入”，恰好对应“ $x$ 是某个祖先的轻儿子的后代”这种情况，次数正好等于 $x$ 到根路径上的轻边条数，不超过 $\log_2 n$ 。

因此总加入次数为 $\sum_x (L(x) + 1) = O(n\log n)$ ，加上每个点一次 $O(1)$ 判定，总复杂度

O(n \log n)

第四步：用 dfs 序把“加入整棵子树”变成区间循环

我们需要一个操作：把整棵子树 $u$ 的所有点一次性加入或移出。由于 $F_i < i$ ，可以用一个手写栈（不需要递归）求出 dfs 序 dfn[] 和 euler[]，此时子树 $u$ 的结点恰好是

\texttt{euler[dfn[u]}\; \dots \;\texttt{dfn[u]} + \mathrm{siz}[u] - 1\texttt{]}

一段连续区间。于是“加入整棵子树”就是一个普通 for 循环，既快又不需要递归。

另外， $F_i < i$ 还让子树大小和重儿子可以在 $O(n)$ 时间内用一次倒序扫描求出来，不必 DFS：倒着枚举 $i = n, n-1, \dots, 2$ ，把 siz[i] 累加到 siz[F_i]，并顺手用 siz[heavy[f]] 的比较维护重儿子。

第五步：把递归也换成迭代

到这里算法已经完整了，但还有一个实现陷阱：DSU on Tree 一般是递归写的，而 $n$ 可以到 $2\times 10^5$ 。当数据是一条链时，递归深度就是 $2\times 10^5$ ，会直接把栈压爆（本机实测段错误）。

解决办法是把递归改成“显式栈 + 阶段标记”。观察递归版 dfs_sack(u, keep) 在每个结点上做三件事：

处理所有轻儿子（各自算完清空）；
处理重儿子（算完保留）；
补上轻儿子和自己、判定答案、按需清空。

于是给每个结点一个阶段 stage[u]：弹出 $u$ 时若 stage[u] == 0，就压入所有轻儿子并把 stage[u] 置 1；若 stage[u] == 1，就压入重儿子并把 stage[u] 置 2；若 stage[u] == 2，就执行第 3 件事并真正出栈。

keep 参数也换成数组：一个结点是父亲的重儿子就保留，否则清空。因为变长部分（儿子列表的遍历）全部用 for 循环完成、不需要入栈，所以每个结点最多同时占用一个栈位。

例如对下面这棵以 1 为根的小树（2 是 1 的重儿子，6 是 2 的重儿子，7 是 3 的重儿子），实际的判定顺序是 $7 \to 3 \to 4 \to 5 \to 6 \to 2 \to 1$ ：

text

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20        1
       / \
      2   3          siz[2] = 4, siz[3] = 2  -> 1 的重儿子是 2
    / | \  \
   4  5  6  7        siz[4]=siz[5]=siz[6]=siz[7]=1
                     -> 2 的重儿子是 6，3 的重儿子是 7

处理顺序与统计状态（{...} 表示当前统计里的点）：
  ① 轻儿子 3：
       先算它的重儿子 7           -> {7}，判定 7，保留
       再补入轻儿子（无）和自己 3  -> {7,3} = S_3，判定 3
       3 不是 1 的重儿子，整体清空 -> {}
  ② 重儿子 2：
       先算轻儿子 4、5（各算完就清）-> {} -> {}
       再算重儿子 6               -> {6}，判定 6，保留
       补入轻儿子 4、5 和自己 2    -> {6,4,5,2} = S_2，判定 2
       2 是 1 的重儿子，保留       -> {6,4,5,2}
  ③ 回到 1：
       此时统计里已经是 S_2，再补入轻儿子 3 的整棵子树 {3,7} 和自己 1
                                  -> {6,4,5,2,7,3,1} = S_1，判定 1

读这张图时注意两点：判定 $u$ 时统计范围总是恰好等于 $S_u$ ；算法对每个点只判定一次，而被反复“加入”的只有轻儿子子树里的点（例子里是结点 3、7 所在的子树被 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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193/**
 * Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
 * rbook: -> https://rbook.roj.ac.cn  https://rbook2.roj.ac.cn
 * rainboy的学习导航网站: https://idx.roj.ac.cn
 * create_at: 2026-10-02 15:07
 * update_at: 2026-10-03 11:07
 */
// main.cpp：P9233 颜色平衡树。
// 判断每个结点的子树是否“每种颜色的点数都相同”，用 DSU on Tree（树上启发式合并）做到 O(n log n)。
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int MAXN = 200005;

int n;
int col[MAXN];              // col[i]：结点 i 的颜色
int fa[MAXN];               // fa[i]：结点 i 的父亲，fa[1] = 0
vector<int> son_list[MAXN]; // son_list[i]：结点 i 的孩子列表

int siz[MAXN];        // siz[u]：子树大小
int heavy[MAXN];      // heavy[u]：u 的重儿子编号，0 表示没有重儿子
bool keep_flag[MAXN]; // keep_flag[u]：算完子树 u 后是否保留统计信息（u 是父亲的重儿子时保留）
int dfn[MAXN];        // dfn[u]：u 的 dfs 序
int euler[MAXN];      // euler[k]：dfs 序为 k 的结点编号
int timer;            // dfs 序计数器
int stk[MAXN];        // 求 dfs 序时用的手写栈
int work_stk[MAXN];   // 模拟 dfs_sack 时用的手写栈
int stage[MAXN];      // stage[u]：结点 u 当前处在 dfs_sack 的哪一步

// 当前统计范围内（一段连续结点）的颜色信息
int cnt[MAXN]; // cnt[c]：颜色 c 出现的次数
int occ[MAXN]; // occ[k]：恰好出现 k 次的颜色种数
int kind;      // 当前统计范围内出现过的颜色种数

ll ans; // 颜色平衡的子树个数

// 把结点 u 的颜色计数加上 delta（delta = +1 或 -1），同时维护 occ 和 kind。
void add_point(int u, int delta) {
    int c = col[u];
    if (cnt[c] > 0) {
        occ[cnt[c]]--; // 颜色 c 的旧次数不再计数
    } else {
        kind++; // 颜色 c 第一次出现
    }
    cnt[c] += delta;
    if (cnt[c] > 0) {
        occ[cnt[c]]++; // 颜色 c 的新次数开始计数
    } else {
        kind--; // 颜色 c 的点全部被移走
    }
}

// 把整棵子树 u 的结点全部加入 / 移出统计范围。
// 子树在 dfs 序上是一段连续的区间，所以这里是一次简单的区间循环。
void add_subtree(int u, int delta) {
    int l = dfn[u];
    int r = dfn[u] + siz[u] - 1;
    for (int i = l; i <= r; i++) {
        add_point(euler[i], delta);
    }
}

// 求子树大小和重儿子。题目保证 fa[i] < i，倒着扫一遍就能把孩子的大小汇总给父亲。
void build_size() {
    for (int i = 1; i <= n; i++) {
        siz[i] = 1;
        heavy[i] = 0;
        keep_flag[i] = false;
    }
    for (int i = n; i >= 2; i--) {
        int f = fa[i];
        siz[f] += siz[i];
        if (siz[i] > siz[heavy[f]]) { // siz[0] 始终为 0，所以 heavy[f] == 0 时也能直接比较
            heavy[f] = i;
        }
    }
    // 重儿子算完后要保留信息给父亲复用，也就是递归版里参数 keep = true 的那些结点。
    keep_flag[1] = true;
    for (int i = 2; i <= n; i++) {
        keep_flag[i] = (heavy[fa[i]] == i);
    }
}

// 手写栈求 dfs 序。把一个点弹出时立刻给它编 dfs 序，再把它的全部孩子压栈。
// 因为栈是后进先出，弹出 u 之后压入的孩子会排在 u 的兄弟之上，
// 所以子树里的点会在 u 之后被连续弹出，每个点的后代恰好构成一段连续区间。
void build_euler() {
    int top = 0;
    stk[++top] = 1;
    timer = 0;
    while (top > 0) {
        int u = stk[top];
        top--;
        timer++;
        dfn[u] = timer;
        euler[timer] = u;
        for (int i = 0; i < (int)son_list[u].size(); i++) {
            stk[++top] = son_list[u][i];
        }
    }
}

// DSU on Tree 主体。
// 递归版每个结点要处理三件事：① 先算所有轻儿子并清空；② 再算重儿子并保留；
// ③ 补上轻儿子和自己，判断答案，最后决定是否清空。
// 这里用 work_stk 模拟这个过程：stage[u] = 0 / 1 / 2 分别表示“还没开始”“轻儿子刚算完”
// “重儿子刚算完”，以便同一个结点被访问三次时做不同的事。
// 用迭代而不用递归，是为了防止链状数据（深度可达 2e5）把递归栈压爆。
void run_sack() {
    int top = 0;
    work_stk[++top] = 1;
    stage[1] = 0;

    while (top > 0) {
        int u = work_stk[top];

        if (stage[u] == 0) {
            // 第一次访问 u：把所有轻儿子压栈，让它们各自单独算并清空。
            stage[u] = 1;
            for (int i = (int)son_list[u].size() - 1; i >= 0; i--) {
                int v = son_list[u][i];
                if (v != heavy[u]) {
                    top++;
                    work_stk[top] = v;
                    stage[v] = 0;
                }
            }
            continue;
        }

        if (stage[u] == 1) {
            // 第二次访问 u：轻儿子都算完并清空了，再算重儿子，作为 u 的“底子”。
            stage[u] = 2;
            if (heavy[u] != 0) {
                top++;
                work_stk[top] = heavy[u];
                stage[heavy[u]] = 0;
            }
            continue;
        }

        // 第三次访问 u：此时统计范围正好是子树重儿子，补上轻儿子子树和 u 自己，
        // 得到的就恰好是整棵子树 u。
        for (int i = 0; i < (int)son_list[u].size(); i++) {
            int v = son_list[u][i];
            if (v != heavy[u]) {
                add_subtree(v, 1);
            }
        }
        add_point(u, 1);

        // 判断子树 u 是否颜色平衡。
        // 设子树大小为 S、出现过的颜色数为 kind。
        // 每种颜色点数都相同 <=> 每种颜色都出现 S / kind 次 <=> kind 整除 S 且 occ[S / kind] == kind。
        if (siz[u] % kind == 0 && occ[siz[u] / kind] == kind) {
            ans++;
        }

        // 轻儿子要把信息清掉；重儿子（含 u 的整棵子树）留给父亲复用。
        if (!keep_flag[u]) {
            add_subtree(u, -1);
        }
        top--;
    }
}

void read_input() {
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> col[i] >> fa[i];
    }
}

void solve() {
    for (int i = 2; i <= n; i++) {
        son_list[fa[i]].push_back(i);
    }
    build_size();
    build_euler();
    run_sack();
    cout << ans << "\n";
}

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

    read_input();
    solve();

    return 0;
}

复杂度

求子树大小与重儿子： $O(n)$ 。
求 dfs 序： $O(n)$ 。
DSU on Tree：每个点被加入 $O(\log n)$ 次，每次 add_point 是 $O(1)$ ，共 $O(n\log n)$ 。
判定：每个点 $O(1)$ ，共 $O(n)$ 。

时间和空间分别是

O(n \log n), \qquad O(n + \max C)

其中 $\max C \leqslant 2\times 10^5$ 。在 $n = 2\times 10^5$ 的链、菊花、完全二叉树、随机树上实测均约 $0.02$ 秒。

总结

这道题的难点不在数据结构，而在于两个需要想清楚的地方：

判定条件的等价改写。 把“所有出现过的颜色次数都相同”改写成“ $k \mid S$ 且恰好有 $k$ 种颜色出现 $S/k$ 次”，从而把判定降到 $O(1)$ 。这一步是所有优化的基础——如果不做这一步，即使统计范围已经就绪，判定本身仍要遍历颜色。
选择性地保留信息。 统计状态只有一份，不可能继承所有儿子，所以要挑一个最有价值的儿子留下。选“子树最大的儿子”，换来“每个点最多被重加 $\log n$ 次”的结论，这就是 DSU on Tree 的全部精髓。

实现层面还有一条经验值得记住：在 $n$ 可达 $2\times 10^5$ 的树上，任何形如“沿子树递归下去”的写法都要警惕栈溢出。这道题里 dfs 序和 DSU on Tree 主体都改成了显式栈，链状数据才不会段错误。

同类问题的识别特征：需要对每个子树回答一个可以用少量全局计数器 $O(1)$ 判定的问题。这时 DSU on Tree 通常比“线段树合并”“树上莫队”更短、更好写、常数更小。