D - Range Count Query Editorial by MMNMM


セグメント木にソート済み列や多重集合を乗せることでもこの問題を解くことができます。

葉に単一の要素 \(A[i]\) のみからなるソート済み列や多重集合を乗せて、葉以外の頂点は子のソート済み列や多重集合を合併したものを乗せることにします。 このとき、ソート済み列や多重集合のサイズの合計は \(N\log N\) と評価できます(深さが同じ頂点のサイズの合計を考えることで示せます)。

このようなセグメント木をもち、クエリごとに対応する頂点のソート済み列や多重集合にそれぞれに対して「\(X\) がいくつ含まれるか」を求め、それらを合計することでこの問題を解くことができました。

この方針では atcoder::segtree をそのまま用いて解くことは難しいため注意してください。

以下の実装では多重集合を \((\)\(, \)個数\()\) の組として連想配列(ハッシュマップ)で持つことで時間計算量を \(\operatorname{expected} O((N+Q)\log N)\) としています。実際の提出

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

int main(){
    // 入力
    int N;
    cin >> N;

    vector<unordered_map<int, int>> segment_tree(2 * N);
    // 葉は {A[i]}
    for(int i = 0; i < N; ++i){
        int A;
        cin >> A;
        segment_tree[i + N][A] = 1;
    }
    // 子を多重集合としてマージ
    for(int i = N; --i;){
        for(const auto [value, count] : segment_tree[2 * i])segment_tree[i][value] += count;
        for(const auto [value, count] : segment_tree[2 * i + 1])segment_tree[i][value] += count;
    }

    int Q;
    cin >> Q;
    for(int i = 0; i < Q; ++i){
        int L, R, X;
        cin >> L >> R >> X;
        L += N - 1;
        R += N;
        int ans = 0;
        // 区間に対応するセグメント木の頂点に対して X の個数を足す
        while(L != R){
            if(L & 1){
                if(segment_tree[L].count(X))ans += segment_tree[L][X];
                ++L;
            }
            if(R & 1){
                --R;
                if(segment_tree[R].count(X))ans += segment_tree[R][X];
            }
            L /= 2;
            R /= 2;
        }
        cout << ans << endl;
    }

    return 0;
}

posted:
last update: