【2023蓝桥杯c++A组】【启发式合并】颜色平衡树

 题目描述：

        给定一棵树，结点由 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）。

难点分析

1. 重儿子的概念与应用：理解并正确应用重儿子的概念是算法的关键。重儿子是指子树大小最大的儿子，通过优先遍历重儿子可以减少递归的深度，从而提高算法的效率。在实际编码中，需要正确地标记每个节点的重儿子，并在DFS遍历中优先处理。

2. 颜色计数的维护：算法中需要维护一个颜色计数数组（cnt）来记录当前子树中每种颜色的数量。同时，为了快速判断当前子树是否满足条件，还需要维护颜色计数的和数组（sum）以及当前子树中颜色的最大和最小值（mx和mn）。在DFS遍历过程中，需要正确地更新这些数组的值。

3. DFS遍历的递归实现：DFS遍历本身是一个递归过程，对于树的算法来说，递归实现是常见的。但是，在本算法中，由于需要处理重儿子和轻儿子的不同情况，并且需要在遍历过程中维护多个数组，因此递归实现的细节较多，容易出错。

4. 边界条件和异常情况的处理：在算法实现过程中，需要特别注意边界条件和异常情况的处理。例如，当树中只有一个节点时，或者当某个节点没有子节点时，需要特别处理。此外，还需要考虑输入数据的有效性验证等问题。

5. 算法复杂度的分析：虽然算法通过重儿子的概念优化了计算过程，但是仍然需要遍历整棵树。因此，算法的时间复杂度为O(n)，其中n为树中节点的数量。在实际应用中，还需要考虑空间复杂度以及常数因子的影响。