【题解】P9233 [蓝桥杯 2023 省 A] 颜色平衡树

P9233 [蓝桥杯 2023 省 A] 颜色平衡树

简要题意

给定一颗 \(n\) 个节点、根为 \(1\) 的树，每个节点有一个颜色 \(c_i\)，问有多少节点满足子树中每一种颜色出现的次数相等。

\(1\le n\le 2\times 10^5\)，\(1\le c_i\le 2\times 10^5\)。

题解

知识点：树，树上启发式合并。

启发：

• 启发就是启发。

• 一种快速处理树上问题的方法。

主题库中最水的 dsu on tree，套个板子改下判定方式就行了，这个算法一句话概括就是保留重儿子子树贡献，暴力删轻儿子子树贡献。

判定方式：

记 \(cnt_i\) 为子树中颜色 \(i\) 的出现次数，\(cc_i\) 为子树中出现次数为 \(i\) 次的颜色的种类数，\(sz_i\) 为节点 \(i\) 子树的大小，\(col_i\) 为节点 \(i\) 的颜色。

那么如果一个点 \(u\) 的子树满足条件，则有 \(\large cnt_{col_u}\times cc_{cnt_{col_u}}=siz_u\)。

#include<bits/stdc++.h>
using namespace std;

#define rep(i,l,r) for(int i=(l);i<=(r);++i)
#define per(i,l,r) for(int i=(r);i>=(l);--i)
#define pr pair<int,int>
#define fi first
#define se second
#define pb push_back
#define all(x) (x).begin(),(x).end()
#define sz(x) (x).size()

#define N 202504
// #define int long long

int n,son[N],siz[N],col[N],cnt[N],cc[N],ans;
vector<int>e[N];

inline void dfs0(int k){
    siz[k]=1;

    for(int x:e[k]){
        dfs0(x);

        siz[k]+=siz[x];

        if(siz[x]>siz[son[k]]){
            son[k]=x;
        }
    }
}

inline void add(int k,int d){
    cc[cnt[col[k]]]--;
    cnt[col[k]]+=d;
    cc[cnt[col[k]]]++;

    for(int x:e[k]){
        add(x,d);
    }
}

inline void calc(int k,bool r){
    for(int x:e[k]){
        if(x==son[k]){
            continue;
        }
        calc(x,0);
    }

    if(son[k]){
        calc(son[k],1);
    }

    cc[cnt[col[k]]]--;
    cnt[col[k]]++;
    cc[cnt[col[k]]]++;

    for(int x:e[k]){
        if(x==son[k]){
            continue;
        }

        add(x,1);
    }

    if(cnt[col[k]]*cc[cnt[col[k]]]==siz[k]){
        ans++;
    }

    if(!r){
        add(k,-1);
    }
}

signed main(){
    // freopen(".in","r",stdin);
    // freopen(".out","w",stdout);
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);

    cin>>n;
    rep(i,1,n){
        int c,f;
        cin>>c>>f;
        e[f].pb(i);
        col[i]=c;
    }

    dfs0(1);
    calc(1,1);

    cout<<ans;

    return 0;
}