题目描述
给定一棵树，结点由 1 至 n 编号，其中结点 1 是树根。树的每个点有一个颜色 Ci。
如果一棵树中存在的每种颜色的结点个数都相同，则我们称它是一棵颜色平衡树。
求出这棵树中有多少个子树是颜色平衡树。输入格式
输入的第一行包含一个整数 n ，表示树的结点数。
接下来 n 行，每行包含两个整数 Ci , Fi，用一个空格分隔，表示第 i 个结点的颜色和父亲结点编号。
特别地，输入数据保证 F1 为 0 ，也即 1 号点没有父亲结点。保证输入数据是一棵树。输出格式
输出一行包含一个整数表示答案。样例输入
6
2 0
2 1
1 2
3 3
3 4
1 4样例输出
4提示
编号为 1, 3, 5, 6 的 4 个结点对应的子树为颜色平衡树。
对于 30% 的评测用例，n ≤ 200，Ci ≤ 200 ；
对于 60% 的评测用例，n ≤ 5000，Ci ≤ 5000 ；
对于所有评测用例，1 ≤ n ≤ 200000，1 ≤ Ci ≤ 200000，0 ≤ Fi < i 。
思路题解
思路：
-
要判断每个子树是否为平衡树，需要统计子树的每种颜色的节点的数量，并判断所有数量是否相等。
-
对于一颗树的根节点，若该树的所有子树的统计结果都得到了，就可以直接将子树的统计结果累加，并加上根节点的颜色。因此可以使用dfs对树进行搜索，在后序遍历位置得到子树的统计结果并累加，就可以计算出该树的统计结果，判断所有颜色数量是否相等即可。
注意：
- 统计结果cnt使用数组时，需要判断整颗树所有颜色的数量，而部分子树的颜色并不包含所有的颜色，每次判断的时间复杂度为O(num_c)，num_c为整棵树的颜色种数，这样会超时。因此可以使用map数据结构，这样每次只需判断子树所包含的颜色。
#include<iostream>
#include<vector>
#include<cstring>
#include<map>
using namespace std;
const int N = 2e5+1;
//最终结果
int ans=0;
//将子树的计数结果cnt_nb累加到根节点的结果cnt上
void add(map<int,int>& cnt,map<int,int>& cnt_nb){
for(auto entry:cnt_nb){
int c=entry.first,count = entry.second;
cnt[c] += count;
}
}
/*
对树进行dfs搜索，树的根节点为i，并返回该子树的各节点颜色计数结果
*/
map<int,int> dfs(vector<int>* g,int* c,int i){
int sz = g[i].size();
map<int,int> cnt; //记录子树的每个节点的各颜色节点的数量
/*如果为叶子节点，直接返回*/
if(sz==0){
cnt[c[i]] = 1;
ans++;
return cnt;
}
/*如果不是叶子节点*/
//将根节点的颜色加入cnt
cnt[c[i]]=1;
//遍历根节点的所有子树，并将子树的计数结果累加到cnt中
for(int j=0;j<sz;j++){
int nb = g[i][j];
map<int,int> cnt_nb = dfs(g,c,nb);
add(cnt,cnt_nb);
}
//判断该子树的各种颜色节点的数量是否相等
int count = cnt[c[i]];
for(auto entry:cnt){
//存在一种颜色数量不等，直接返回
if(entry.second != count) return cnt;
}
//各颜色的数量相等，结果+1
ans++;
//返回计数结果
return cnt;
}
int main()
{
int n;
cin>>n;
vector<int> g[N];
int c[N]; //每个节点的颜色
for(int i=0;i<n;i++){
int f;
cin>>c[i]>>f;
if(f>=1){
g[f-1].push_back(i); //记录节点f的子节点i（节点编号从0开始）
}
}
dfs(g,c,0);
cout<<ans;
return 0;
}