Luogu P5901 [IOI2009] Regions

Hanx16Kira · 2023-12-16 21:18:07 · 题解

[IOI2009] Regions

Luogu P5901

Solution

先尝试去想暴力，很容易想到两种不同的暴力做法：

• 分别枚举两个颜色中的所有点，利用 Dfs 序判定一个点是否在另一个点的子树内。
• 将一个颜色内的所有点先加入到一个数据结构内，然后枚举另一个颜色的所有点，查看有多少点在当前子树。

第一种做法时间最劣是 \mathcal O(n^2) 的，第二种做法空间最劣是 \mathcal O(n^2) 的。考虑综合一下这两种做法。使用根号分治，按照颜色的点数进行根号分治。将点数在阈值之上的称为重颜色，否则称为轻颜色。分类讨论一下：

• 重颜色作为祖先节点：不妨预处理所有重颜色作为祖先节点的情况的答案，容易做到 \mathcal O(n\sqrt {n\log n}) 时间预处理，\mathcal O(R\sqrt n) 空间。
• 轻颜色作为祖先节点：枚举轻颜色内的所有点，然后对于每个颜色开一个 vector 按照 Dfs 序将所有点排序就可以直接二分找到所有满足条件的子节点，时间可以做到 \mathcal O(n\sqrt{n\log n})。

总时间复杂度 \mathcal O(n\sqrt {n\log n})，空间复杂度 \mathcal O(R\sqrt n)。

#include <bits/stdc++.h>
#define All(x) x.begin(), x.end()
using namespace std;
constexpr int _N = 2e5 + 5, _B = 450 + 5, _R = 2.5e4 + 5;
int N, R, Q, Lim;
vector<int> e[_N];
int f[_B][_R], dfn[_N], cnt[_N], ord, id[_N], siz[_N], di[_N], d[_N];
vector<int> vec[_N], vdfn[_N];
void Dfs(int x) {
    dfn[x] = ++ord, siz[x] = 1;
    for (int v: e[x]) Dfs(v), siz[x] += siz[v];
}
signed main() {
    cin >> N >> R >> Q;
    Lim = sqrt(N * log2(N) * 2);
    for (int i = 1; i <= N; ++i) {
        if (i != 1) {
            int fa; cin >> fa;
            e[fa].push_back(i);
        }
        int col; cin >> col;
        vec[col].push_back(i);
        ++cnt[col];
    }
    iota(id + 1, id + R + 1, 1);
    sort(id + 1, id + R + 1, [&](const int &i, const int &j) {
        return cnt[i] > cnt[j];
    });
    Dfs(1);
    for (int i = 1; i <= R; ++i) {
        di[id[i]] = i;
        for (int x: vec[i]) vdfn[i].push_back(dfn[x]);
        sort(All(vdfn[i]));
    }
    for (int i = 1; i <= R; ++i) {
        if (cnt[id[i]] < Lim) break;
        for (int i = 1; i <= N + 1; ++i) d[i] = 0;
        for (int x: vec[id[i]]) ++d[dfn[x]], --d[dfn[x] + siz[x]];
        partial_sum(d, d + N + 1, d);
        for (int j = 1; j <= R; ++j) for (int x: vec[j])
            f[i][j] += d[dfn[x]];
    }
    while (Q--) {
        int r1, r2; cin >> r1 >> r2;
        if (cnt[r1] < Lim) {
            int res = 0;
            for (int x: vec[r1])
                res += upper_bound(All(vdfn[r2]), dfn[x] + siz[x] - 1) -
                       lower_bound(All(vdfn[r2]), dfn[x]);
            cout << res << endl;
        } else cout << f[di[r1]][r2] << endl;
    }
}