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

// 存图方式使用前向星

const int N = 2e5 + 5;

int h[N], ne[N], e[N], idx; // 前向星
int root[N]; // 每一个节点的权值线段树的根节点

int merge(int p, int q, int l, int r) // p 和 q 表示合并的两棵树的根节点，l 和 r 表示当前处理的范围
{
    // ...
}

int calc(int u) // 该函数有两个作用: 1.计算 u 的子树的答案; 2.合并 u 的子树
{
    int res = 0;
    
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i]; // u 为当前节点，j 为正在处理的儿子节点
        res += calc(j); // 先递归处理儿子节点
        
        root[u] = merge(root[u], root[j], 1, 2e5); // 合并线段树
    }
    
    update(res, tr[root[u]].ans); // 因为当前节点的信息已经和所有儿子的信息合并，且所有儿子的信息也和儿子的儿子的信息合并，所以 root[u] 为根的权值线段树的信息为子树 u 的信息
    return res;
}

int merge(int p, int q, int l, int r) // p 和 q 表示合并的两棵树的根节点，l 和 r 表示当前处理的范围
{
    if (!p || !q) return p + q; // 如果不全有节点，那么直接返回有的节点
    if (l == r) 
    {
        tr[p].min = tr[p].max = tr[p].max + tr[q].max; // 因为是叶子节点，所以直接相加
        return p;
    }
    
    int mid = l + r >> 1;
    tr[p].l = merge(tr[p].l, tr[q].l, l, mid), tr[p].r = merge(tr[p].r, tr[q].r, mid + 1, r); // 递归的合并左右儿子
    tr[p].max = max(tr[tr[p].l].max, tr[tr[p].r].max);
    tr[p].min = min(tr[tr[p].l].min, tr[tr[p].r].min); // 上传标记
    return p;
}

void update(int &res, int root)
{
    res += (tr[root].min == tr[root].max); // 当最小值与最大值相等，每个值都相等，产生贡献
}

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

#define int long long
const int N = 2e5 + 5, inf = 1e18;
int h[N], ne[N], e[N], idx;
int col[N], root[N], n;

void add(int a, int b)
{
    ne[idx] = h[a], e[idx] = b, h[a] = idx ++;
}

struct Node {
    int l, r, min, max;
} tr[N << 5];

int tot;

int createNode()
{
    int cur = ++ tot;
    tr[cur] = {0, 0, inf, 0};
    return cur;
}

void pushup(int u)
{
    tr[u].min = min(tr[tr[u].l].min, tr[tr[u].r].min);
    tr[u].max = max(tr[tr[u].l].max, tr[tr[u].r].max);
}

void insert(int &p, int l, int r, int pos, int delta)
{
    if (!p) p = createNode();
    if (l == r) 
    {
        tr[p].max = tr[p].min = tr[p].max + delta;
        return ;
    }

    int mid = l + r >> 1;
    if (pos <= mid) insert(tr[p].l, l, mid, pos, delta);
    else insert(tr[p].r, mid + 1, r, pos, delta);
    pushup(p);
}

int merge(int p, int q, int l, int r)
{
    if (!p || !q) return p + q;
    if (l == r)
    {
        tr[p].max += tr[q].max;
        tr[p].min += tr[q].min;
        return p;
    }
    int mid = l + r >> 1;

    tr[p].l = merge(tr[p].l, tr[q].l, l, mid), tr[p].r = merge(tr[p].r, tr[q].r, mid + 1, r);
    pushup(p);
    return p;
}

int dfs(int u)
{
    int res = 0;
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        res += dfs(j);
        merge(root[u], root[j], 1, 2e5);
    }

    return res + (tr[root[u]].max == tr[root[u]].min);
}

int32_t main()
{
    tr[0].min = inf;
    memset(h, -1, sizeof h);
    
    cin >> n;
    for (int i = 1; i <= n; ++i)
    {
        int fa;
        cin >> col[i] >> fa;
        if (fa) add(fa, i);
        insert(root[i], 1, 2e5, col[i], 1);
    }

    cout << dfs(1) << '\n';

    return 0;
}