区間 k-最小値いろいろ

ここで区間 \(k\)-最小値と呼んでいるのは以下のような問題です。

最初に配列 \(a = (a_0, a_1, \dots, a_{n-1})\) が与えられるよ。 次にクエリ処理を \(q\) 回してね。 クエリは \( ([l, r), k)\) の形式で与えられるので、\(a_l, a_{l+1}, \dots, a_{r-1}\) を昇順ソートしたときに 0-indexed で先頭から \(k\) 番目に来る値を答えてね。 ただし \(0\le k \lt r-l\) が保証されるので、答えがなくなることはないから安心してね。

昨日の ABC-D は、これに \(l = 0\) の追加制約が与えられたものと見ることができます。

座標圧縮などをすれば帰着できるので、以下では \(0\le a_i \lt n\) とします。値の重複も \( (a_i, i)\) などで tie-break することにすればよいので、distinct であるとします（つまり、\(0, 1, \dots, n-1\) の順列が与えられるものとします*1）。

この記事では次のことを紹介します。

• wavelet matrix での解法（の概要）
  • 前処理 \(O(n\log(n))\) time、クエリ \(O(\log(n))\) time
  • クエリ先読みがいらない
    • \(i\) 個目のクエリに答えたら \(i+1\) 個目のクエリが与えられる、などの状況でもいける
• 並列二分探索での解法
  • バッチ処理 \(O( (n+q)\log(n)^2)\) time
  • クエリ先読み

並列二分探索の方が計算量では劣っているのですが、並列二分探索の導入や実装の練習（あるいはライブラリの verify）として使えそうだな〜と思ったので、書こうと思い至った次第です。 自分の中ではこのアプローチにさっき初めて思い至った（気がする）のと、あまり言及されているのを見た記憶がない（気がする）ので。

二つの解法は各々別物なので、片方だけ読みたいなら他方は飛ばして問題ないはずです。

wavelet matrix での解法

wavelet matrix というデータ構造があります。以下では WM と書きます。

競プロ界隈ではセグ木ほどの知名度はない印象ですが、整数に関するいろいろな区間クエリを処理できて強力です。 たとえば、区間 \(k\)-最小値クエリなどを処理できるので、そのまま使って終わりです。

ここでは WM 自体の詳細には立ち入らないですが、このクエリを処理する際に WM がしていることの概要を説明します。

まず、\(n \lt 2^m\) なる最小の \(m\) を求めておきます。 答えは \(b = b_{m-1}\dots b_1 b_0\) (\(b_i \in \{0, 1\}\)) の形で表せます。

さて、\(j\gets m-1, \dots, 1, 0\) に対して順に以下の処理をします。

• 区間 \([l, r)\) にある要素で上位 bit が \(b_{m-1}\dots b_{j+1}\) であるもののうち、\(j\) 番目の bit が \(0\) であるものの個数 \(c\) を数える。
• \(c\gt k\) であれば、\(b_j = 0\) で確定する。
  • 下位 bit の並び方には影響されず、注目中の bit で決まるので。
  • 注目中の bit に \(0\) が十分個あれば \(k\)-最小値の該当 bit が \(0\) とわかるということ。
• \(c\le k\) であれば、\(b_j = 1\) で確定する。
  • 小さい方から \(c\) 個は注目中の bit が \(0\) なので、注目中の bit が \(1\) のもののうち、\( (k-c)\)-最小値を求めればいいとわかる。
    • \(k\gets k-c\) で更新する。

以上により、\(b\) の全 bit が確定します。

よって、各 \(j\) について、「\(j\) 番目の bit が \(0\) であるものの個数の区間和クエリに答えられるデータ構造を、各 \(b_{m-1}\dots b_{j+1}\) ごとに持ったもの」があれば十分とわかります。 これは（愚直には）\(b_{m-1}\dots b_{j+1}\) でソートして、その値ごとにグルーピングして各々累積和を管理するみたいなことをすれば実現できる気がします*2。

詳細は他に任せます。

並列二分探索での解法

まず並列二分探索のフレームワークを説明してから、区間 \(k\)-最小値への適用のしかたを説明します。

並列二分探索自体を知っている人は前半は斜め読みでいいと思います。

並列二分探索のフレームワーク

以下のような問題を考えます。

機械があり、初めは状態 \(0\) です。状態 \(i\) (\(0\le i\lt m\)) のとき、状態 \(i+1\) に遷移させることができます。 また、いつでも状態 \(0\) に戻すことができます。

ある \(0\le i^{\star}\le m+1\) が存在して、状態 \(i\lt i^{\star}\) のとき、機械はあなたの質問に対して true と返します。 状態 \(i\ge i^{\star}\) のときは false と返します。

常に true を返すケースや、常に false を返すケースも存在することに注意してください。

あなたは、なるべく少ない質問回数で \(i^{\star}\) を特定したいです。

二分探索すればよいです。

let mut machine = Machine::new();
if !machine.query() { return 0; }

let mut yes = 0;
let mut no = m + 1;
while no - yes > 1 {
    let mid = yes + (no - yes) / 2; // 0 <= mid <= m
    machine.reset_state();
    for _ in 0..mid {
        machine.next_state();
    } // 状態 mid になっている
    *(if machine.query() { &mut yes } else { &mut no }) = mid;
}
no // no == i_star

これでは、質問ごとに状態をリセットして mid まで遷移させ直しているのがもったいなく見えます。 状態の遷移を \(O(m\log(m))\) 回と、質問を \(O(\log(m))\) 回しています。

そこで、質問の種類を増やしてみます。

機械があり、初めは状態 \(0\) です。状態 \(i\) (\(0\le i\lt m\)) のとき、状態 \(i+1\) に遷移させることができます。 また、いつでも状態 \(0\) に戻すことができます。

ある \(0\le i^{\star}_j\le m+1\) が存在して、状態 \(i\lt i^{\star}_j\) のとき、機械はあなたの質問 \(j\) (\(0\le j\lt n)\) に対して true と返します。 状態 \(i\ge i^{\star}_j\) のときは false と返します。

常に true を返すケースや、常に false を返すケースも存在することに注意してください。

あなたは、なるべく少ない質問回数で各 \(i^{\star}_j\) を特定したいです。

各 \(j\) について「状態なにのときに質問するか」を求め、機械がその状態になったら質問 \(j\) をすることにします。 これにより、一度状態 \(0\) から状態 \(m\) に遷移させる間に \(n\) 回質問できます。 \(n\) 個の質問たちが必要とする遷移たちが共通なので、まとめられるわけですね。

これを \(\log(m)\) 回繰り返せば各 \(i^{\star}_j\) を特定できるので、全体で \(O(m \log(m))\) 回の遷移と \(O(n\log(m))\) 回の質問をしたことになります。 質問が \(n\) 個になったので質問回数は元々の \(n\) 倍になっていますが、遷移回数は悪化していません。

並列二分探索の適用

並列二分探索で \( (0, 1, \dots, n-1)\) の順列 \(a = (a_0, a_1, \dots, a_{n-1})\) の区間 \(k\)-最小値問題を解くためのやり方です。

上で言っていた機械を作ります。 ここでは、機械は以下の配列 \(b = (b_0, b_1, \dots, b_{n-1})\) を管理します。配列は BIT かセグ木の形で持ちます。

• 状態 \(0\) においては、\(b\) の各要素は \(0\) とする。
• 状態 \(i\) (\(0\le i\lt n\))から状態 \(i+1\) に遷移するとき、\(a_j = i\) なる \(j\) に対して \(b_j \gets 1\) とする。
  • \(a\) を値の小さい順に使っていくイメージ。
  • 転倒数や LIS を求めるときと同じ、いわゆる平面走査。
• 質問 \( ([l, r), k)\) に対しては、\( (\sum_{j\in [l, r)} b_i) \le k\) を返す。
  • 小さい順に値を入れていき、\(k+1\) 個目に入ったものが \(k\)-最小値であることに注意。

質問に対する true/false は、区間 \(k\)-最小値が未確定かどうかに対応しています。 状態 \(i+1\) で初めて false になったのであれば、該当のクエリの答えは \(i\) ということになります。 今回の状況設定では「常に true」「常に false」という事態は発生しません。

全体で \(O(n \log(n))\) 回の遷移と \(O(q\log(n))\) 回の質問で、各々一回あたり \(O(\log(n))\) 時間なので、全体では \(O( (n+q)\log(n)^2)\) 時間です。

実装例

• library-checker/range_kth_smallest
• atcoder/abc234_d

コンテスト中の提出ですが、WM を用いた提出 も一応貼っておきます。

あとがき

長くなってしまいました。

計算量などの兼ね合いからあまり使うことはなさそうですが、「セグ木しかないが区間 \(k\)-最小値を解く必要が出てきた」といった状況*3では役に立つことと思います。

「log 倍悪くなってしまうが手持ちのデータ構造を使えば解ける」といったものは、ライブラリの verify で役に立つことがそれなりにあるので、そうした観点では面白い気もします。

区間を BTreeSet<(int, int)> で管理するやつの verify に、区間代入クエリが使える話なども 過去に書いた のですが、これもお気に入りです。

おわり

おつかれさまですにゃ。

*1:\( (\max_i a_i) \ll n\) の場合などはそれを利用して計算量が落ちたりしますが、ここでは割愛します。

*2:実際の WM では、こんな感じのことを暗に表現しています。実装もそこまで重くはないです。

*3:ある？