KotlinのBinarySearchを二分探索に使いたい場合の挙動が微妙にむずかしい
KotlinのArrayやListにはbinarySearch()というメソッドが生えているので、ソート済みのリストに対して二分探索による高速な値の検索が簡単に行うことが出来ます。
(JavaにはArrays.binarySearch()というメソッドがあり、おそらく内部的には同じもの)
さて、このメソッドは単にリスト内に存在する値のindexを返すというだけの用途であれば特に迷うことがないですが、競技プログラミングをやっていると「対象の値未満のもっとも近いindex」などを調べたいケースがままあります。
これはC++などではlower_boundやupper_boundといった関数が用意されていますが、Kotlinではどうやら存在しないようなので、(関数を自作しないのであれば)binarySearchを利用する必要があります。
リストに値がない場合にbinarySearchメソッドが返す値は少し直感的ではありません。
リファレンスによると
Return the index of the element, if it is contained in the list within the specified range; otherwise, the inverted insertion point (-insertion point - 1).
値が存在しない場合はその値をlistにsort列を満たすように挿入する時の最小インデックスをXとするとき,-X-1を返す。
と書かれているようです。
ちょっとわかりにくいですね。
実際に挙動を確かめてみました。
val list = listOf(10, 20, 30, 40, 50)
debugLog(list.binarySearch(10)) // 0 対象の値が存在するのでindexを返す。
debugLog(list.binarySearch(11)) // -2 対象の値が存在しない。対象の値より大きい最も近い値のindex + 1を負数で返す
debugLog(list.binarySearch(19)) // -2 対象の値が存在しない。対象の値より大きい最も近い値のindex + 1を負数で返す
debugLog(list.binarySearch(29)) // -3 対象の値が存在しない。対象の値より大きい最も近い値のindex + 1を負数で返す
debugLog(list.binarySearch(5)) // -1 最低値未満でも同じ
debugLog(list.binarySearch(55)) // -6 最大値以上でも同じ。この場合正の数にした時にsizeよりも大きい値が返る
まず、対象の値が存在しない場合は負数が返ってきます。
そしてその数は正の数に直すと対象の値より大きい最も近い値のindex + 1になっています。
上の例では、11を検索した時に直感的に期待する値は1ですが、11は存在しないため負数が返ります。更に10のindexである0 + 1の負数である-1ではなく、-2が返ってきます。
-X-1なので、正の数に直すと+1されているように見えるのですね。
対象の値未満のもっとも近いindexを求めたい
たとえばリストが
[10,20,30,40,50]
で、対象の値が25だった場合、25未満でもっとも大きな数のindexを探したいですね。
この場合、期待する値は20のindexである1です。
しかし、list.binarySearch(25)の結果は-3になります。
そのため、
-list.binarySearch(25) - 2
とすることで、期待する値である1を取得することができます。
対象の値より大きい最も近いindexを求めたい
逆パターンです。
対象の値が25だった場合、期待する値は30のindexである2です。
この場合も実際には-3が返ってくるため、
-list.binarySearch(25) - 1
とすればよいです。
対象の値が含まれる場合
ここまでリスト内に対象の値そのものが含まれない前提で考えてきましたが、値が見つかってしまうと上記のコードはバグってしまいます。
なのでまず結果が正の値か負の値かを判定し、処理を分ける必要があります。
対象の値そのものが検索にヒットしても問題がないというケースの場合は単に正負の判定をするだけですが、対象の値は無視してもっとも近い値のindexを見つけたいケースもあります。
この場合、binarySearchの引数で独自のComparatorクラスを渡して同じ値を無視する必要があります。
※参考
たとえば20の位置を見つけたい場合、以下のようなコードの結果は意図通りではありません。
val list = listOf(10, 20, 25, 25, 25, 30, 40, 50)
debugLog(list.binarySearch(25)) // 3
これに上の記事で紹介されているLowerBoundComparatorを使うことで、25が配列に含まれない場合と同じく-3が返ってくるようになります。
list.binarySearch(25, LowerBoundComparator()) // -3
逆に30の位置を見つけたい場合は、期待する値は5です。
今度はUpperBoundComparatorを利用すると期待する値を見つけることが出来ます。
list.binarySearch(25, UpperBoundComparator()) // -6
おもしろいですね。(白目)