【颜色平衡树 / E】

题目

思路

• DFS暴力 60分

代码

#include <bits/stdc++.h>
using namespace std;
const int N = 5010;
const int M = 5010;
int h[N], e[M], ne[M], idx;
int c[N], f;
int ans;
void add(int a, int b)  // 添加一条边a->b
{
    e[idx] = b, ne[idx] = h[a], h[a] = idx ++ ;
}
bool check(int tmp[])
{
    int last = 0;
    for(int i = 1; i <= 5000; i++)
    {
        if(!last && tmp[i]) last = tmp[i];
        else if(last && tmp[i])
        {
            if(last != tmp[i]) return false;
            last = tmp[i];
        }
    }

    return true;
}
void dfs(int u, int f[])
{
    int tmp[N] = {0};
    //根节点颜色算进子树
    tmp[c[u]] += 1;
    
    bool leaf = true;
    for(int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        dfs(j, tmp);
        leaf = false;
    }
    
    //判断平衡性
    if(leaf) ans++;
    else
    {
        if(check(tmp)) ans++;
    }
    
    //呈递颜色
    for(int i = 1; i <= 5000; i++)
        f[i] += tmp[i];
}
int main()
{
    int n;
    cin >> n;
    
    memset(h, -1, sizeof h);
    for(int i = 1; i <= n; i++)
    {
        cin >> c[i] >> f;
        add(f, i);
    }
    
    int tmp[N] = {0};
    dfs(1, tmp);
    cout << ans;
}