はまやんはまやんはまやん

hamayanhamayan's blog

Range Count Query [ユニークビジョンプログラミングコンテスト2022(AtCoder Beginner Contest 248) D]

https://atcoder.jp/contests/abc248/tasks/abc248_d

前提知識

解説

https://atcoder.jp/contests/abc248/submissions/31044881

色々な解法が見える問題。
実際WaveletMatrixというデータ構造を使えば、一瞬で答えが出るのだが、二分探索解法で説明する。

以下、配列の要素については0-indexed(0始まり)で説明している。

手動でクエリを計算してみる

手動で答えてくださいと言われたら、L番目からR番目まで順番にXがあるかを見ていくような感じになると思う。
ここで、Xだけを見ることが出来れば、高速化できそうという気持ちになり、そういった方法を探してみる。

数が同じものをまとめておく

クエリを処理前の段階でとある数について、その数がある添え字(index)をまとめた配列を用意しておくことにする。

indexes[x] := 値がxである数列の要素の添え字の集合

具体的には入力例1 3 1 4 1 5であれば、

indexes[1] = { 1, 3 }
indexes[3] = { 0 }
indexes[4] = { 2 }
indexes[5] = { 4 }

というのをあらかじめ用意しておく。
こうすれば、とある数Xが含まれる場所は前計算で求めたこれを使えば高速に求められることになる。

要求問題が変わる

こうすると、要求される問題が以下のように変化する。

与えられた配列の中にLからRの数は何個あるか

これを計算するには二分探索が使える。
二分探索するためには配列がソートされている必要があるので注意。
自分の作成方法だと、構築すると既にソートされているが、実装によってはそうでないかもしれない。
自分の作成方法はコードを確認してほしい。

さて、二分探索で求めるのは、

  • L以上の数が初めて現れる場所 st
  • Rより大きい数が初めて現れる場所 gl

である。これが求まれば、個数はgl - stとなるので、これを求めればいい。

C++であれば、iteratorというのを理解する必要があるが、

  • ソート済みの配列に対して~以上の数が初めて現れる場所はlower_bound
  • ソート済みの配列に対して~より大きいの数が初めて現れる場所はupper_bound

で二分探索を意識しなくても計算することができる。
二分探索ではO(logN)の計算量なので、全体でO(QlogN)となって間に合う。

int N, A[201010];
int Q;
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N) cin >> A[i];
    cin >> Q;

    map<int, vector<int>> indexes;
    rep(i, 0, N) indexes[A[i]].push_back(i);

    rep(q, 0, Q) {
        int L, R, X;
        cin >> L >> R >> X;
        L--; R--;

        auto st = lower_bound(all(indexes[X]), L);
        auto gl = upper_bound(all(indexes[X]), R);

        int ans = gl - st;
        printf("%d\n", ans);
    }
}