公式

E - 会社の組織と給与 / Company Organization and Salaries 解説 by MMNMM


まず、オイラーツアーを利用して頂点を並べ替えることで、部分木に対するクエリを区間に対するクエリに読み替えることができます。

この読み替えによって、解くべき問題は次のような形になります。

長さ \(N\) の数列 \(A=(A _ 1,A _ 2,\ldots,A _ N)\) がある。次の \(2\) 種類のクエリを \(Q\) 個処理せよ。

  1. \(i,x\) が与えられ、\(A _ i\) を \(x\) に更新する。
  2. \(l,r,v\) が与えられ、\(A _ l,A _ {l+1},\ldots,A _ r\) のうち \(v\) 以上であるものの個数を求める。

この問題は、クエリ先読みをすることでセグメント木や BIT のそれぞれの頂点に座標圧縮をしたセグメント木や BIT を持って解くことができます。

具体的には以下のようなものを考えることで解くことができます。

セグメント木のうち、区間 \(\lbrack L,R\rparen\) を保持する頂点は

  • 集合 \(S\) を \(\bigl\lbrace\)クエリを \(k\) 個 \((0\le k\le Q)\) 処理した直後の \(A _ i\ (L\le i\lt R)\) の値\(\bigr\rbrace\) として定め、これを \(S=\lbrace s _ 1,s _ 2,\ldots,s _ n\rbrace\ (s _ 1\lt s _ 2\lt\cdots\lt s _ n)\) とする。 \(T\) の区間 \(\lbrack l,r\rparen\) への取得クエリの結果が、「\(i=L,L+1,\ldots,R\) のうち \(s _ l\le A _ i\lt s _ r\) を満たすものの個数」となる

ようなセグメント木 \(T\) を持つ。

時間計算量は \(O((N+Q)\log N\log (N+Q))\)、空間計算量は \(O((N+Q)\log N)\) などとなります。

実装例は以下のようになります。

#include <iostream>
#include <vector>
#include <algorithm>
#include <atcoder/fenwicktree>
using namespace std;

int main() {
    int N, Q;
    cin >> N >> Q;

    vector<int> A(N);
    for (int& a : A) {
        cin >> a;
    }

    vector<vector<int>> children(N);
    for (int i = 1; i < N; ++i) {
        int P;
        cin >> P;
        --P;
        children[P].emplace_back(i);
    }

    // DFS でオイラーツアーを求める
    vector<pair<int, int>> euler_tour(N);
    [&children, &euler_tour](this auto dfs, int now, int count) -> int {
        euler_tour[now].first = count;
        ++count;
        for (int c : children[now]) {
            count = dfs(c, count);
        }
        euler_tour[now].second = count;
        return count;
    }(0, 0);

    // クエリを先読み
    vector<tuple<int, int, int>> queries(Q);
    for (auto& [t, v, x] : queries) {
        cin >> t;
        if (t == 1) {
            cin >> v >> x;
            --v;
        } else {
            cin >> v;
            --v;
        }
    }

    // 座標圧縮を行う
    vector<vector<int>> compress(2 * N);
    // ありえる値をすべて追加
    for (int i = 0; i < N; ++i) {
        compress[euler_tour[i].first + N].emplace_back(A[i]);
    }
    for (auto [t, v, x] : queries) {
        if (t == 1) {
            compress[euler_tour[v].first + N].emplace_back(x);
        }
    }
    // ソートして重複を削除する
    for (int i = 0; i < N; ++i) {
        ranges::sort(compress[i + N]);
        auto removed = ranges::unique(compress[i + N]);
        compress[i + N].erase(removed.begin(), removed.end());
    }
    for (int i = N; --i > 0;) {
        ranges::copy(compress[2 * i], back_inserter(compress[i]));
        ranges::copy(compress[2 * i + 1], back_inserter(compress[i]));
        ranges::sort(compress[i]);
        auto removed = ranges::unique(compress[i]);
        compress[i].erase(removed.begin(), removed.end());
    }

    // 各頂点に BIT を持つセグメント木を作成
    vector<atcoder::fenwick_tree<int>> bits;
    for (int i = 0; i < 2 * N; ++i) {
        bits.emplace_back(atcoder::fenwick_tree<int>(size(compress[i])));
    }

    // 単一の BIT に対する操作を定義しておく
    //     更新: x の個数に v を足す
    auto add_bit = [&compress, &bits](int i, int x, int v) {
        int bit_index = ranges::lower_bound(compress[i], x) - begin(compress[i]); // 座圧を復元し
        bits[i].add(bit_index, v); // BIT を更新
    };
    //     取得: x より大きい個数を求める
    auto get_bit = [&compress, &bits](int i, int x) {
        int bit_index = ranges::upper_bound(compress[i], x) - begin(compress[i]); // 座圧を復元し
        return bits[i].sum(bit_index, size(compress[i])); // BIT から取得
    };

    // 全体に対する操作
    //      頂点 v の値を x に変更
    auto update = [N, &A, &euler_tour, &add_bit](int v, int x) {
        int index = euler_tour[v].first + N;
        while (index) {
            add_bit(index, A[v], -1); // A[v] を 1 減らして
            add_bit(index, x, 1); // x を 1 増やす
            index /= 2;
        }
        A[v] = x; // A を更新
    };
    //     頂点 v の子のうち v より大きい個数を求める
    auto get = [N, &A, &euler_tour, &get_bit](int v) {
        int l = euler_tour[v].first + N, r = euler_tour[v].second + N;
        int ret = 0;
        while (l < r) {
            if (l & 1) {
                ret += get_bit(l++, A[v]);
            }
            if (r & 1) {
                ret += get_bit(--r, A[v]);
            }
            l /= 2;
            r /= 2;
        }
        return ret;
    };

    // 初期値を追加
    for (int i = 0; i < N; ++i) {
        int index = euler_tour[i].first + N;
        while (index) {
            add_bit(index, A[i], 1);
            index /= 2;
        }
    }

    // クエリを順に処理
    for (auto [t, v, x] : queries) {
        if (t == 1) {
            update(v, x);
        } else {
            cout << get(v) << endl;
        }
    }
    return 0;
}