昇順にソートされたリストから任意の値のインデックスを取得したいことがある。
lower_bound()を使えばいいかもしれないけど、検索対象のデータ構造がvector<pair<int, int>>で、検索したいのがペアの1つ目の要素であるようなとき、lower_bound()を使う方法だと少しよくわからない記述をしないといけない。
他にも競プロでは「lower_bound()が使えない/使いづらいケース」がけっこうあって、そのたびに自前で二分探索を書くのだけど、そのたびに少し悩むのでルーチン化した。
まず、検索対象の数列の要素を横に並べる。
次に欲しいインデックスの場所に矢印を書き加える。今回は5の右端が欲しいので下の図の矢印の位置に書き加える。そしてこの矢印にOKという名前を付ける。
次にこの矢印のひとつ左にNGの矢印を加える。
最後に、「OKの矢印の要素が満たし、NGの矢印の要素が満たさないような条件」を考える。すると、OKの要素の条件は「5を超える」ことだとわかる。ということで、二分探索のOK、NGの更新式は次のようになる。
(x > 5 ? ok : ng) = md;
まとめると、次のような手順を踏むと二分探索を書く準備が整う。
- 数列を書く
- 欲しいインデックスの場所にOKの矢印を加える
- そのひとつ左にNGの矢印を加える
- OKの要素が満たし、NGの要素が満たさないような条件を考える
OKとNGの初期位置は次のようにする。
先ほど書いた「OKの矢印の要素」というのは、OKの矢印の右側にある要素のことを指すので、初期位置は矢印の右側の要素が空になるようにする。
もし「OKの矢印の要素」が左側の要素を指すとしたら、初期位置は次のように変わる。
OKの矢印が右側の2つの要素を指す場合は、次のような初期位置になる。
つまり、矢印の参照する範囲の場所がわかれば、矢印の初期位置は自ずと決まる。
そして、参照する範囲と矢印の位置関係はルールとして決めておくと楽である。僕の場合は、矢印が範囲の左端を指すようにと決めている。もし右端としてしまうと、要素を参照する際にA[i - 1]のように- 1を書かなければならないからである。
左端と決めると、NGの矢印の場所も「OKの矢印のひとつ左」というように自然に決まってくる（なぜ決まるかについてはうまく説明できない）。すると条件も自ずと求まる。
結論
ということで、考えなければならないのは「OKの矢印をどこに書くか」だけである。これさえきちんと決めることができれば、NGの矢印の位置、条件式、矢印の初期位置、これらは芋づる式に決まっていく。
例題1
問題
5以下の要素の個数を求めたいとき、下図のどこにOKの矢印を加えればいいか。
解答
以下の場所に加えればいい。もし要素番号が0-indexであれば、OKの値がそのまま5以下の要素の個数となる。
繰り返しになるが、OKの位置が決まることで、NGの位置、条件式、矢印の初期位置が芋づる式に決まる。今回の場合では、NGの位置はOKのひとつ左、条件式は「要素の値が5を超えるか」、初期位置は、この記事の一番最初に載せた初期位置の図の位置である。
例題2
問題
5以上の要素の個数を求めたいとき、下図のどこにOKの矢印を加えればいいか。
解答
以下の場所に加えればいい。もし要素番号が0-indexであって、数列の長さがNのとき、NからOKの値を引いたものが解となる。
条件式は「要素の値が5以上」となる。NGの矢印の位置と初期位置は例題1と同じ。
まとめ
OKの矢印をどこに書くかだけ考えればいい。そのために、数列を紙に書き出す作業をすると楽になる。
