蓝桥杯2023年第十四届省赛真题E 颜色平衡树

#include <bits/stdc++.h>
using namespace std;
int ans = 0;
map<int, int>mp;//分别表示颜色编号，颜色的个数

//节点对应颜色的数组
//节点对应儿子的vector
//遍历就好了
int color[200001] = {0}; //下标表示节点编号，值表示颜色编号
vector<int> son[200001];//下标表示节点编号，
//vector表示这个节点的儿子节点编号

void dfs(int i, map<int, int> &mp) { //节点编号，表示
	if (son[i].size() == 0) { //这个节点是末端节点
		mp[color[i]]++;//先把自己的信息加进去
		ans++;
		return;
	}
	//如果不是叶子节点，遍历其所有儿子节点，先找出其所有儿子节点
	//for(auto j:mp)
	map<int, int>a; //前一个表示颜色编号，后一个表示个数
	a[color[i]]++;//建立子树的颜色个数map后，先把自己加进去
	for (auto j : son[i])
		dfs(j, a);

	int pre = 0, f = 1;
	for (auto j : a) { //遍历子树的颜色个数map
		if (pre != 0 && pre != j.second) {
			f = 0;
		}
		mp[j.first] += j.second;
		pre = j.second;
	}
	if (f == 1)
		ans++;
}

int main() {
	int sum;//总节点数
	cin >> sum;
	int c, f;
	for (int i = 1; i <= sum; i++) {
		cin >> c >> f;
		color[i] = c;
		son[f].push_back(i);
	}
	dfs(1, mp);
	cout << ans;
	return 0;
}