2023第十四届蓝桥杯省赛C++A组：颜色平衡树（词典记录颜色+递归解法

题目描述
给定一棵树，结点由 1 至 n 编号，其中结点 1 是树根。树的每个点有一个颜色 Ci。
如果一棵树中存在的每种颜色的结点个数都相同，则我们称它是一棵颜色平衡树。
求出这棵树中有多少个子树是颜色平衡树。

输入格式
输入的第一行包含一个整数 n，表示树的结点数。
接下来 n 行，每行包含两个整数 Ci, Fi，用一个空格分隔，表示第 i 个结点的颜色和父亲结点编号。
特别地，输入数据保证 F1 为 0 ，也即 1 号点没有父亲结点。
保证输入数据是一棵树。
对于 30% 的评测用例，n≤200，Ci≤200 ；
对于 60% 的评测用例，n≤5000，Ci≤5000 ；
对于100% 的评测用例，1≤n≤200000，1≤Ci≤200000，0≤Fi<i 。

输出格式
输出一行包含一个整数表示答案。

解题思路

对每一个子树判断是否为颜色平衡树，先考虑递归解法，故dfs用来判断该结点作为父节点的树是否为颜色平衡树，并用哈希表来记录各种颜色的数量。

对于叶子节点，因其只有一个颜色，故一定为颜色平衡树，以此作为递归终止条件。用tree容器记录各个节点的孩子，用color记录各个节点对应的颜色，在进行递归时，遍历孩子，并将各子树的颜色数进行求和，然后利用其中一个颜色的数量来判断该节点对应的子树是否为颜色平衡树

附代码+注释

#include<bits/stdc++.h>
#include<unordered_map>
using namespace std;
const int N = 200000;
vector<vector<int>> tree(N + 1);
vector<int> color(N+1);
int sum = 0;
unordered_map<int, int> dfs(int node, unordered_map<int, int>& map) {
	if (tree[node].empty())
	{
		sum++;                      //叶子节点必为颜色平衡树
		map[color[node]]++;
		return map;
	}
	else {
		for (auto x : tree[node])
		{
			unordered_map<int, int> it;
			dfs(x, it);
			for (auto m : it) {
				map[m.first] += m.second;     //将子树的颜色求和
			}
		}
		map[color[node]]++;
		int count = map[color[node]];
		for (auto m : map) {
			if (m.second != count)
				return map;					  //判断是否为颜色平衡树
		}
		sum++; return map;
	}
}
int main() {
	int n;
	cin >> n;
	int p = 1;
	unordered_map<int, int> map;
	while (n--) {
		int c;
		int temp;
		cin >> c >> temp;
		color[p] = c;                    //储存颜色
		tree[temp].push_back(p);         //记录节点孩子
		p++;
	}
	dfs(1, map);
	cout << sum;
	return 0;
}