P9233 [蓝桥杯 2023 省 A] 颜色平衡树（dfs序莫队）

莫队原理:https://zhuanlan.zhihu.com/p/115243708

对于树上的每个结点，按照 dfs 序打上时间戳，这样就可以把每一个结点对应的子树的答案转化为一个区间的答案。将子树询问离线下来变成 $n$ 个区间查询。

用 $cnt$ 数组维护颜色的出现次数，那么如果一个区间有贡献，当且仅当 $cnt$ 数组中所有非零元素都一样，表示每种颜色的节点个数相同。用 $sum$ 维护 $cnt$ 中有多少种互不相同的元素，那么当 $sum$ 等于 $1$ 时，说明每种颜色的出现次数均相同，颜色平衡树的个数加一。

我们还需要用一个数组 $tot$ 维护当前 $cnt$ 中每种元素的出现次数，如果 $tot$ 的某一个元素从 $0$ 变为 $1$ ，相当于 $cnt$ 中多了一种元素，将 $sum$ 加一；如果 $tot$ 中的某一个元素从 $1$ 变为 $0$ ，相当于 $cnt$ 数组中少了一种元素，将 $sum$ 减一。

#include<bits/stdc++.h>

using namespace std;
const int N = 2e5 + 10;
int n, m, x, t, num, sum, ans, c[N];
int in[N], out[N], id[N], pos[N], cnt[N], tot[N];
vector<int> g[N];

struct Node {
    int l, r;
    bool operator <(const Node &t) const {
        if (pos[l] != pos[t.l]) return l < t.l;
        if (pos[l] & 1) return r < t.r;
        return r > t.r;
    }
} q[N];

void dfs(int x) {
    in[x] = ++num, id[num] = x;
    for (int y : g[x])    dfs(y);
    out[x] = num;
    q[++t] = {in[x], out[x]};
}

void add(int x) {
    if (cnt[c[x]]) {
        if (--tot[cnt[c[x]]] == 0)    sum--;
    }
    cnt[c[x]]++;
    if (++tot[cnt[c[x]]] == 1)    sum++;
}

void del(int x) {
    if (--tot[cnt[c[x]]] == 0)    sum--;
    cnt[c[x]]--;
    if (cnt[c[x]]) {
        if (++tot[cnt[c[x]]] == 1)    sum++;
    }
}

int main() {
    cin >> n;
    m = sqrt(n);
    for (int i = 1; i <= n; i++) {
        pos[i] = (i - 1) / m + 1; //分块
        cin >> c[i] >> x;
        g[x].push_back(i);
    }
    dfs(1), sort(q + 1, q + 1 + n);
    for (int l = 1, r = 0, i = 1; i <= n; i++) {
        while (l > q[i].l) add(id[--l]);
        while (r < q[i].r) add(id[++r]);
        while (l < q[i].l) del(id[l++]);
        while (r > q[i].r) del(id[r--]);
        if (sum == 1)    ans++;
    }
    cout << ans << endl;
}