题目描述:
给定一棵树,结点由 1 至 n 编号,其中结点 1 是树根。树的每个点有一个颜色 Ci。
如果一棵树中存在的每种颜色的结点个数都相同,则我们称它是一棵颜色平衡树。
求出这棵树中有多少个子树是颜色平衡树。
题目链接:
P2071 - [蓝桥杯2023初赛] 颜色平衡树 - New Online Judge
#include<bits/stdc++.h>
using namespace std;
using pii = pair<int, int>;
const int N = 2e5;
// 主函数
int main() {
ios::sync_with_stdio(false); // 关闭C++与C的输入输出同步,以加快输入输出速度
cin.tie(0); // 将cin与cout的绑定取消,以加快输入输出速度
int n;
cin >> n; // 读取树的节点数量
// 初始化颜色数组和邻接表
vector<int> color(n + 1);
vector<vector<int>> g(n + 1);
// 读取每个节点的颜色和父节点,并构建邻接表
for(int i = 1; i <= n; i++) {
int fa;
cin >> color[i] >> fa;
if (fa != 0) {
g[fa].push_back(i);
g[i].push_back(fa); // 这里是无向图,所以两个方向都要添加
}
}
// 初始化各种辅助数组
vector<int> cnt(N + 1), L(n + 1), R(n + 1), big(n + 1), sz(n + 1), node(n + 1), sum(n + 1);
int dfn = 0; // 深度优先搜索的编号
int ans = 0, mx = 0, mn = 1e9; // ans是答案,mx和mn分别记录当前子树中某种颜色的最大和最小数量
// 第一次深度优先搜索,计算每个节点的子树大小、深度优先搜索编号和重儿子(子树最大的儿子)
function<void(int, int)> dfs1 = [&](int u, int fa) {
sz[u] = 1; // 初始化当前节点的子树大小为1
L[u] = ++dfn; // 分配深度优先搜索编号
node[dfn] = u; // 记录下该编号对应的节点
// 遍历所有儿子节点
for (auto v : g[u]) {
if (v != fa) { // 排除父节点
dfs1(v, u); // 递归搜索儿子节点
sz[u] += sz[v]; // 累加子树大小
// 更新重儿子
if (!big[u] || sz[big[u]] < sz[v]) {
big[u] = v;
}
}
}
R[u] = dfn; // 当前节点的深度优先搜索编号结束位置
};
// 添加节点u的颜色到统计中,并更新mx和mn
auto add = [&](int u) {
sum[cnt[color[u]]]--;
if(sum[cnt[color[u]]] == 0 && mn == cnt[color[u]]) {
mn = cnt[color[u]] + 1;
} else {
mn = min(mn, cnt[color[u]] + 1);
}
cnt[color[u]]++;
sum[cnt[color[u]]]++;
mx = max(mx, cnt[color[u]]);
};
// 清除节点u的颜色统计
auto del = [&](int u) {
mx = 0, mn = 1e9;
sum[cnt[color[u]]] = 0;
cnt[color[u]] = 0;
};
// 第二次深度优先搜索,计算答案
function<void(int, int, bool)> dfs2 = [&](int u, int fa, bool keep) {
for(auto v : g[u]) {
if (v != fa && v != big[u]) {
dfs2(v, u, false);
}
}
// 遍历轻儿子(非重儿子的儿子)
for(auto v : g[u]) {
if (v != fa && v != big[u]) {
dfs2(v, u, false); // 轻儿子不保留统计信息
}
}
// 遍历重儿子(如果有的话)
if (big[u]) {
dfs2(big[u], u, true); // 重儿子保留统计信息
}
for (auto v : g[u]) {
if (v != fa && v != big[u]) {
for (int i = L[v]; i <= R[v]; i++) {
add(node[i]);
}
}
}
// 添加当前节点的颜色到统计中
add(u);
// 如果当前子树中某种颜色的数量最大值和最小值相等,说明所有颜色数量都相同,答案加1
if (mx == mn) {
ans++;
}
// 如果不是重儿子,则清除当前子树的所有颜色统计
if (keep == false) {
for (int i = L[u]; i <= R[u]; i++) {
del(node[i]);
}
dfs1(1, -1);
dfs2(1, -1, false);
cout << ans << "\n";
return 0;
}
小结
本段代码实现了一个基于树的算法,具体为在给定一棵树上,通过两次深度优先搜索(DFS)遍历来计算满足特定条件的子树数量。首先,通过第一次DFS(dfs1
)遍历,计算了每个节点的子树大小、深度优先搜索的编号(L[u]
和R[u]
),并标记了每个节点的重儿子(big[u]
)。然后,在第二次DFS(dfs2
)遍历中,利用了重儿子的性质来优化计算过程,同时通过维护一个颜色的计数数组(cnt
)、颜色计数的和数组(sum
)以及当前子树中颜色的最大和最小值(mx
和mn
)来统计满足条件的子树数量。
算法的核心在于通过重儿子的概念来减少不必要的重复计算,同时利用颜色计数数组来快速判断当前子树是否满足条件(即所有颜色的数量都相等)。最终,算法返回满足条件的子树数量(ans
)。
难点分析
-
重儿子的概念与应用:理解并正确应用重儿子的概念是算法的关键。重儿子是指子树大小最大的儿子,通过优先遍历重儿子可以减少递归的深度,从而提高算法的效率。在实际编码中,需要正确地标记每个节点的重儿子,并在DFS遍历中优先处理。
-
颜色计数的维护:算法中需要维护一个颜色计数数组(
cnt
)来记录当前子树中每种颜色的数量。同时,为了快速判断当前子树是否满足条件,还需要维护颜色计数的和数组(sum
)以及当前子树中颜色的最大和最小值(mx
和mn
)。在DFS遍历过程中,需要正确地更新这些数组的值。 -
DFS遍历的递归实现:DFS遍历本身是一个递归过程,对于树的算法来说,递归实现是常见的。但是,在本算法中,由于需要处理重儿子和轻儿子的不同情况,并且需要在遍历过程中维护多个数组,因此递归实现的细节较多,容易出错。
-
边界条件和异常情况的处理:在算法实现过程中,需要特别注意边界条件和异常情况的处理。例如,当树中只有一个节点时,或者当某个节点没有子节点时,需要特别处理。此外,还需要考虑输入数据的有效性验证等问题。
-
算法复杂度的分析:虽然算法通过重儿子的概念优化了计算过程,但是仍然需要遍历整棵树。因此,算法的时间复杂度为O(n),其中n为树中节点的数量。在实际应用中,还需要考虑空间复杂度以及常数因子的影响。