【Java】lower_bound(), upper_bound() 関数を作る

　「プログラミングコンテストチャレンジブック [第2版]」（通称：アリ本）を見ていて、ちょっと気になる解法があったのだが、C++ のライブラリの「lower_bound()」を使っていたので Java では試せなかった。というわけで Java に lower_bound(), upper_bound() を移植してみた。ちなみにこの２つの関数は二分探索の関数で、lower_bound() は「指定した値以上の先頭の要素を返す」関数で、upper_bound() は「指定した値より大きい先頭の要素を返す」関数だ。Java 標準ライブラリの「Arrays.binarySearch()」も二分探索のライブラリだが、同じ値が並んでいるときは、返ってくる要素は不定になり、見つからなかった場合は負の値になる。これら関数は、基本的に見つかった値に対して、返ってくる要素（インデクス）が違うだけと考えて良い。今回は配列内の要素での探索に使う。

●Java 版 lower_bound(), upper_bound()

/**<h1>指定した値以上の先頭のインデクスを返す</h1>
 * <p>配列要素が０のときは、０が返る。</p>
 * @param arr ： 探索対象配列(単調増加であること)
 * @param value ： 探索する値
 * @return<b>int</b> ： 探索した値以上で、先頭になるインデクス
 */
public static final int lowerBound(final int[] arr, final int value) {
    int low = 0;
    int high = arr.length;
    int mid;
    while (low < high) {
        mid = ((high - low) >>> 1) + low;    //(low + high) / 2 (オーバーフロー対策)
        if (arr[mid] < value) {
            low = mid + 1;
        } else {
            high = mid;
        }
    }
    return low;
}

/**<h1>指定した値より大きい先頭のインデクスを返す</h1>
 * <p>配列要素が０のときは、０が返る。</p>
 * @param arr ： 探索対象配列(単調増加であること)
 * @param value ： 探索する値
 * @return<b>int</b> ： 探索した値より上で、先頭になるインデクス
 */
public static final int upperBound(final int[] arr, final int value) {
    int low = 0;
    int high = arr.length;
    int mid;
    while (low < high) {
        mid = ((high - low) >>> 1) + low;    //(low + high) / 2 (オーバーフロー対策)
        if (arr[mid] <= value) {
            low = mid + 1;
        } else {
            high = mid;
        }
    }
    return low;
}

　気になった解法というのは、「アリ本 [第2版]」の P.135 ページにある「Subsequence」(部分列問題)のはじめの方の解法だ。問題としては、

長さ N の数列 A と、整数 S が与えられます。A の連続する部分列で、その総和が S 以上となるようなもののうち、最小の長さを求めなさい。

というものだが、「最初から各インデクスまでの部分和を配列に記録し、各インデクスからの部分和で S 以上となる、先頭のインデクスを二分探索する」というアルゴリズムで解いている。これは要素がすべて正の数なら、その部分和は単調増加になるので、二分探索が適用できるというものだ。上の「lowerBound()」を使って、Java で書いてみよう。

●Java 版 二分探索を使った Subsequence 解法

//メインコード
int[] A = new int[]{5, 1, 3, 5, 10, 7, 4, 9, 2, 8};    //数列
int N = A.length;    //数列の長さ
int S = 15;        //求める和(以上)

static final int INF = 100000000;

System.out.println(solve());    //答え


//二分探索での解法
int solve() {
    int ans = INF;
    int[] sum = new int[N + 1];

    for (int i = 0; i < N; i++) {
        sum[i + 1] = sum[i] + A[i];
    }

    for (int i = 0; S <= sum[N] - sum[i]; i++) {
        int t = lowerBound(sum, sum[i] + S);
        ans = Math.min(ans, t - i);
    }

    return (ans == INF ? 0 : ans);
}

2

　といっても、この問題はその後にある「しゃくとり法」の方が、より効率が良いんだけどね。ただ、配列の中身がバラバラの数値（単調増加でない数列）だったので、部分和の二分探索は思い付かなかった。単純な全探索でも解けるけど、要素数が多いと計算量が間に合わないからね。解法は色々知っていて損はない。まぁ、ライブラリがないとこの解法はすぐにはできないが…。


　また、これら２つの関数を組み合わせると、「同じ値の連続した区間」を高速に求めることができる。例えば以下のように配列にデータが入っていたとして、探す値を v としたとき、２つの関数の差分を取れば良い。

●「同じ値の連続した区間」を二分探索で求める

int[] data = {1, 2, 2, 3, 3, 3, 4, 5, 5, 6};     //ソート済みであること
int v = 3;    //検索する値
int k = upperBound(data, v) - lowerBound(data, v);    //差分が個数になる(ないとき=0)
System.out.println(k);

3

　これは 3 の値が「連続で 3 個」あることを表している。


　ちなみにもう一つ工夫して、引数に Comparator を使ったりすると、単調減少でも利用することができる。ただし返る値は「見つかった値より後ろのインデクス」のようになるけどね。ほんの少し改造版も載せておこう。

●Comparator を追加した、単調減少で探索できる lower_bound(), upper_bound()

/**<h1>指定した値と同じか後ろにある、先頭のインデクスを返す</h1>
 * <p>配列要素が０のときは、０が返る。</p>
 * @param arr ： 探索対象配列(単調増加/減少であること)
 * @param value ： 探索する値
 * @param comparator ： Comparator<Integer> インターフェイス
 * @return<b>int</b> ： 探索した値と同じか後ろで、先頭になるインデクス
 */
public static final int lowerBound(final int arr[], final int value, final Comparator<Integer> comparator) {
    int low = 0;
    int high = arr.length;
    int mid;
    int cmp;
    while (low < high) {
        mid = ((high - low) >>> 1) + low;    //(low + high) / 2 (オーバーフロー対策)
        cmp = comparator.compare(arr[mid], value);
        if (cmp < 0) {
            low = mid + 1;
        } else {
            high = mid;
        }
    }
    return low;
}

/**<h1>指定した値より後ろにある、先頭の要素を返す</h1>
 * <p>配列要素が０のときは、０が返る。</p>
 * @param arr ： 探索対象配列(単調増加/減少であること)
 * @param value ： 探索する値
 * @param comparator ： Comparator<Integer> インターフェイス
 * @return<b>int</b> ： 探索した値より後ろで、先頭になるインデクス
 */
public static final int upperBound(final int arr[], final int value, final Comparator<Integer> comparator) {
    int low = 0;
    int high = arr.length;
    int mid;
    int cmp;
    while (low < high) {
        mid = ((high - low) >>> 1) + low;    //(low + high) / 2 (オーバーフロー対策)
        cmp = comparator.compare(arr[mid], value);
        if (cmp <= 0) {
            low = mid + 1;
        } else {
            high = mid;
        }
    }
    return low;
}


//メインコード例（A は単調減少の配列）
int ans = lowerBound(A, K, new Comparator<Integer>() {
    @Override
    public int compare(Integer o1, Integer o2) {
        return o2 - o1;    //降順
    }
});

　あとは、探索対象の開始と終了インデクスも引数に追加して、配列の一部のみの探索も可能にするといいね。この記事でのコードは１つ１つの関数で書いてしまっているけど、一番多く引数を持つ関数を１つだけ作って、あとは引数を変えた関数をオーバーロードをしてしまう方が楽だ。必要ならば自分で作ってみると良いだろう。


（関連記事）
【Java】二分探索を汎用的に使えるようにする
【Ruby】二分探索, lower_bound, upper_bound
【Java】配列要素の反転(reverse)
【Java】２次元配列のソート
【Java】連想配列(Map)を値でソートする

■参考になる書籍