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

· · 题解

题目传送门

思路

先考虑暴力。显然对每颗子树做一遍 dfs 判定，如何判定？可以把 C_i 装到桶 \alpha_i 里，设 \beta_i 表示有多少个 \alpha_j 为 i，typ 表示元素种类，siz_i 表示以 i 为根的子树的大小。显然我们只需判定 typ|siz_u 且 \beta_{siz_u/typ}=typ 即可。上述的数组插入 / 删除一个数均能 O(1) 维护，时间复杂度 O(n^2)。

使用 dsu on tree 优化。我们不期望重儿子的子树遍历多次，因为它的节点数最多，损耗的时间最大。于是我们可以把每次重儿子遍历后的状态都保存不动，然后遍历非重儿子的子树，储存计算完后再删掉。因为轻边边数为 \log n 条，所以总复杂度为 O(n\log n)。

Code

以下设 a=C，b=\alpha，c=\beta。

#include <bits/stdc++.h>
using namespace std;
#define il inline
#define pb push_back
const int N = 2e5 + 10;
vector<int> g[N];
int a[N], b[N], c[N], typ, ans, n;
int dp[N], siz[N], son[N], dfn[N], rdfn[N], dfscnt;
il void add(int x)
{
    c[b[x]]--, c[++b[x]]++;
    if(b[x] == 1) typ++;
}
il void del(int x)
{
    c[b[x]]--, c[--b[x]]++;
    if(!b[x]) typ--; 
}
il void dfs1(int u)
{
    dfn[u] = dp[u] = ++dfscnt;
    rdfn[dfscnt] = u;
    siz[u] = 1;
    son[u] = -1;
    for(auto &v : g[u])
    {
        dfs1(v);
        if(son[u] == -1 || siz[v] > siz[son[u]])
            son[u] = v;
        siz[u] += siz[v];
        dp[u] = max(dp[u], dp[v]);
    }
}
il void dfs2(int u, bool tag)
{
    for(auto &v : g[u])
        if(v != son[u])
            dfs2(v, 0);
    if(son[u] != -1)
        dfs2(son[u], 1);
    for(auto &v : g[u])
    {
        if(v != son[u])
        {
            for(int i = dfn[v];i <= dp[v];++i)
                add(a[rdfn[i]]);
        }
    }
    add(a[u]);
    if(siz[u] % typ == 0 && c[siz[u] / typ] == typ)
        ans++;
    if(!tag)
    {
        for(int i = dfn[u];i <= dp[u];++i)
            del(a[rdfn[i]]);
    }
}
int main()
{
    cin >> n;
    for(int i = 1;i <= n;++i)
    {
        int x;
        scanf("%d%d", &a[i], &x);
        if(!x) continue;
        g[x].pb(i);
    }
    c[0] = n;
    dfs1(1);
    dfs2(1, 0);
    cout << ans;
    return 0;
}