蓝桥杯2023年-颜色平衡树（哈希表，dfs，map作返参）

题目描述

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