题目描述

给定一棵树，结点由 1 至 n 编号，其中结点 1 是树根。树的每个点有一个颜色 Ci。

如果一棵树中存在的每种颜色的结点个数都相同，则我们称它是一棵颜色平衡树。

求出这棵树中有多少个子树是颜色平衡树。

思路

（map作返参的用法感觉挺少用的，题目倒不难想。虽然当时也是抱着试一试的态度写的。）

由于颜色的数值是离散的，我们可以使用map记录各颜色数量，聚合一下。然后dfs，先算出所有子树的颜色，然后加上当前根结点的颜色，再判断所有种类的颜色数量是否相等。如果数量都相等，则ans++。然后把当前树的map颜色记录作为返回值返回就行。

代码

#include<bits/stdc++.h>
using namespace std;
vector<int>son[200005];
int color[200005];
int ans=0;
map<int,int>dfs(int u){map<int,int> mp;map<int,int>temp;for(auto it:son[u]){temp=dfs(it);for(auto it1:temp){mp[it1.first]+=temp[it1.first];}}mp[color[u]]++;int k=-1;int flag=1;for(auto it:mp){if(k==-1){k=it.second;}else {if(k!=it.second){flag=0;break;}}}ans+=flag;return mp;
}
int main()
{int n;cin>>n;for(int i=1;i<=n;i++){int a,p;cin>>a>>p;color[i]=a;son[p].push_back(i);}dfs(1);cout<<ans;
}