2023年第十四届蓝桥杯C/C++ A组 E题 颜色平衡树

试题 E：颜色平衡树

要判断每个子树是否为颜色平衡树，需要统计子树各种颜色的节点数量，并判断所有数量是否相等。

对于一棵树的根节点，如果它所有子树的统计结果都得到了，可以将所有子树的统计结果累加并算上根节点的颜色，这样就可以判断是否为颜色平衡树。

因此，我们可以使用深度优先搜索，将后序遍历得到的子树统计结果累加，然后判断颜色是否相等。

考虑访问的便捷性，我们采用 map 存储各个颜色对应的数量。

#include <iostream>
#include <vector>
#include <map>

using namespace std;

int ans = 0; // 最终结果

/**
 * 将子树的颜色数量表添加到根的颜色数量表
 *
 * @param cnt 根的颜色数量表
 * @param cnt_nb 子树的颜色数量表
 */
void add(map<int, int> &cnt, map<int, int> &cnt_nb) {
    for (auto entry: cnt_nb) {
        int color = entry.first;
        int count = entry.second;
        cnt[color] += count;
    }
}

/**
 * 深度优先遍历
 *
 * @param children 孩子数组
 * @param colors 颜色数组
 * @param root 根节点下标
 * @return 以root为根节点的树的各个颜色个数
 */
map<int, int> dfs(vector<vector<int>> &children, vector<int> &colors, int root) {
    int size = children[root].size();
    map<int, int> cnt; // 记录以root为根的树中所有节点的颜色数量
    if (size == 0) {
        cnt[colors[root]] = 1;
        ans++;
        return cnt;
    }
    cnt[colors[root]] = 1;
    // 对每个孩子节点进行遍历
    for (int i = 0; i < size; i++) {
        // 递归寻找每个子树的颜色数量表
        map<int, int> cnt_nb = dfs(children, colors, children[root][i]);
        // 合并
        add(cnt, cnt_nb);
    }
    // 判断是否满足颜色平衡
    int count = cnt[colors[root]];
    for (auto entry: cnt) {
        if (entry.second != count) return cnt;
    }
    ans++;
    return cnt;
}

int main() {
    int n;
    cin >> n;
    vector<int> colors(n); // 颜色数组
    vector<vector<int>> children(n); // 孩子数组
    for (int i = 0; i < n; ++i) {
        int fa;
        cin >> colors[i] >> fa;
        // 节点i是节点fa的孩子
        if (fa > 0) {
            children[fa - 1].push_back(i);
        }
    }
    dfs(children, colors, 0);
    cout << ans << endl;
    return 0;
}