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

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;
}