颜色平衡树

题目描述

https://www.dotcpp.com/oj/problem3144.html

代码

#include<iostream>
#include<vector>
#include<cstdio>
#include<unordered_map>
using namespace std;
 
//color是为了记录对应编号的节点的颜色
//childNode是为了记录对应编号的子节点的编号
//colorcount是每一个节点的对应的颜色和个数之间的映射
//sum是每一个节点的对应的个数和颜色种类的映射
const int N = 200010;
int color[N] , n;
int ans = 0;
vector<int> childNode[N];
unordered_map<int, int> colorcount[N], sum[N];
 
void dfs(int node) {
    colorcount[node][color[node]] = sum[node][1] = 1;
    for(auto& v : childNode[node]) {
        dfs(v);
        for(auto& it : colorcount[v]) {
            int vcolor = it.first, vcolorcount = it.second;
            if(colorcount[node].count(vcolor) && !--sum[node][colorcount[node][vcolor]]){
            
                sum[node].erase(colorcount[node][vcolor]);
                sum[node][colorcount[node][vcolor] += vcolorcount]++;
            }else {
                sum[node][colorcount[node][vcolor] += vcolorcount]++;
            }
        }
    }
    ans += sum[node].size() == 1;
}
int main()
{
    cin>>n;
    for(int i = 1, x; i <= n; i++) {
        scanf("%d%d", color + i, &x);
        childNode[x].push_back(i);
    }
     
    dfs(1);
    cout << ans;
    return 0;
}