java.util.Arrays.binarySearch で二分探索 + Java で std::lower_bound/upper_bound

諸事情により Java を使わなきゃいけなくなった、そして二分探索がやりたくなったのでメモ。

java.util.Collections.binarySearch でもほぼ同様（二分探索の対象が配列か List のようなコレクションかの違い）。

解説

java.util.Arrays.binarySearch は int 型の値を返します。ドキュメント (Arrays (Java Platform SE 6)) 曰く、

Returns:
index of the search key, if it is contained in the array; otherwise, (-(insertion point) - 1). The insertion point is defined as the point at which the key would be inserted into the array: the index of the first element greater than the key, or a.length if all elements in the array are less than the specified key. Note that this guarantees that the return value will be >= 0 if and only if the key is found.

適当に和訳すると、

戻り値：
配列が検索キーを含むならそのインデックスを返す。含まなければ、 (-(挿入位置) - 1) を返す。挿入位置は、「キーをその配列に挿入する場合の位置」と定義される。それはキーより大きい最初の要素のインデックスであるか、その配列の length である（配列のすべての要素が指定されたキーより小さい場合）。これは「キーが見つかった時には戻り値は0以上となり、見つからなかった時には戻り値は負となる」ことを保証する、ということに注意。

つまり、戻り値が負ならビット否定すれば挿入位置を得られるということですね。

コード例

int[] xs = {1, 3, 5};
System.out.println("xs: " + Arrays.toString(xs));

for (int key = 0; key <= 6; ++key) {
    int result = Arrays.binarySearch(xs, key);
    int insertionPoint = (result >= 0) ? result : ~result;
    System.out.printf("key: %d | result: %2d | insertion point: %d\n", key, result, insertionPoint);
}

出力

出力は以下の通り（検索キー、binarySearch の戻り値、挿入位置の順）。

xs: [1, 3, 5]
key: 0 | result: -1 | insertion point: 0
key: 1 | result:  0 | insertion point: 0
key: 2 | result: -2 | insertion point: 1
key: 3 | result:  1 | insertion point: 1
key: 4 | result: -3 | insertion point: 2
key: 5 | result:  2 | insertion point: 2
key: 6 | result: -4 | insertion point: 3

おまけ：std::lower_bound/upper_bound 風 Comparator

配列が指定キーを複数含む場合、先頭の方にある要素のインデックスが返ってくるとは限らない（要素が1個見つかった時点で探索は打ち切られる）ので、C++ の std::lower_bound（指定した値以上の要素が初めて出現する場所を取得）や std::upper_bound（指定した値より大きい要素が初めて出現する場所を取得）のようには使えないということに注意。

以下のような「等価であると判断しない（compare が 0 を返さない）Comparator」を定義し、インスタンスを binarySearch に引数として渡せばそれっぽいことができる。

import java.util.Comparator;

class LowerBoundComparator<T extends Comparable<? super T>>
    implements Comparator<T>
{
    public int compare(T x, T y)
    {
        return (x.compareTo(y) >= 0) ? 1 : -1;
    }
}

class UpperBoundComparator<T extends Comparable<? super T>>
    implements Comparator<T>
{
    public int compare(T x, T y)
    {
        return (x.compareTo(y) > 0) ? 1 : -1;
    }
}

ただし、ジェネリクスは int などのプリミティブな型には適用できず、そもそも int[] 用の binarySearch は Comparator を引数に取ってくれないため、int[] を Integer[] に変換する必要があったりして大変ダルい模様。

• 参考: java - How to sort an array of ints using a custom comparator? - Stack Overflow
  • ベストアンサーでは Apache Commons の ArrayUtils で int[] と Integer[] を相互に変換している

利用例は以下の通り（ここでは java.util.List と java.util.Collections.binarySearch を利用）。

List<Integer> xs = Arrays.asList(1, 2, 2, 3, 3, 3);
System.out.println("xs: " + xs);

Map<String, Comparator<Integer>> map = new LinkedHashMap<>();
map.put("Lower Bound", new LowerBoundComparator<>());
map.put("Upper Bound", new UpperBoundComparator<>());

for (Map.Entry<String, Comparator<Integer>> e : map.entrySet()) {
    System.out.println("-- " + e.getKey() + " --");
    Comparator<Integer> comparator = e.getValue();
    for (int key = 0; key <= 4; ++key) {
        int result = Collections.binarySearch(xs, key, comparator);
        int insertionPoint = (result >= 0) ? result : ~result;
        System.out.printf("key: %d | result: %2d | insertion point: %d\n", key, result, insertionPoint);
    }
}

出力は以下の通り（検索キー、binarySearch の戻り値、挿入位置の順）。

xs: [1, 2, 2, 3, 3, 3]
-- Lower Bound --
key: 0 | result: -1 | insertion point: 0
key: 1 | result: -1 | insertion point: 0
key: 2 | result: -2 | insertion point: 1
key: 3 | result: -4 | insertion point: 3
key: 4 | result: -7 | insertion point: 6
-- Upper Bound --
key: 0 | result: -1 | insertion point: 0
key: 1 | result: -2 | insertion point: 1
key: 2 | result: -4 | insertion point: 3
key: 3 | result: -7 | insertion point: 6
key: 4 | result: -7 | insertion point: 6