pyてよn日記

一寸先は闇が人生

競プロ覚書：二分探索，std::lower_bound を使いこなす

AtCoder C++ std::lower_bound 二分探索 競プロ覚書

螺旋本で二分探索を勉強したのでその内容をまとめました．二分探索を「使いこなす」などと大きいことを言いましたが，基礎的な競プロの問題を解いているだけです（ごめんなさい）．

結構長めなのでお時間があるときに気になったところからでも読んでいただけたらありがたいです．

概要

大まかに以下の流れで解説していきます．

  1. 二分探索の基礎
  2. C++ で二分探索の実装
  3. C++ 標準ライブラリで二分探索を行う
  4. 標準ライブラリの活用
  5. 標準ライブラリを用いて競プロの問題を解く

記事がだいぶ長くなってしまったので，ライブラリの使いたい方を知りたいだけの方は 3 以降を見るだけで十分だと思います．問題を解きたい方は 5 に飛んでください．

二分探索とは

二分探索は，データの大小関係を利用した高速な探索アルゴリズムです．ソートされている配列に対して高速に要素の検索を行います．

アルゴリズム

二分探索のアルゴリズムは以下になります．ソートされた配列に対して適用します（ここ大事）

  1. 配列全体を探索の範囲とする．
  2. 探索の範囲内の中央に位置する要素を調べる．
  3. 目的のキー（検索対象）と中央の要素が一致すれば探索を終了．
  4. 目的のキーが見つからない場合
    • キー < 中央の要素：探索範囲を前半部分に絞って 2 に戻る
    • 中央の要素 < キー：探索範囲を後半部分に絞って 2 に戻る

以下の画像でイメージが掴めると思います．

f:id:pytwbf201830:20190220191628p:plain

（AOJ：二分探索 http://judge.u-aizu.ac.jp/onlinejudge/commentary.jsp?id=ALDS1_4_B&pattern=post&type=general&filter=Algorithm より引用）

各計算ステップが終わるごとに調べる範囲が半分になっていくため，端から検索していく線形探索よりも高速に探索を行うことができます．特に，ソートをすれば二分探索が使えるという考え方は覚えておくべきだと思います．

計算量

二分探索の計算量は O(\log N) です．これは，一回の比較演算ごとに探索の範囲が半分になっていくことから容易に求められます．

どれくらい効率的かというのは実際の演算回数を見れば直感的に分かると思います．以下に，要素数 N の配列に対して線形探索，二分探索を適用した場合の最悪の演算回数（最悪計算量）を表にしてみました．

要素数 N 線形探索 二分探索
100 100 回 7 回
10,000 10,000 回 14 回
1,000,000 1,000,000 回 20 回

このように比較すると，二分探索が以下に効率的か良く分かると思います．

一つ注意点としては，二分探索を行う際に，配列をソートするための計算量も同時に考えなければいけません．高速なソートアルゴリズムの計算量は O(N \log N) なので，これも念頭に置いて二分探索を行いましょう．

C++ で二分探索を実装

二分探索の基礎を確認したところで，C++ を使って実装していきます．二分探索の実装にはインデックスの取り方などにより，色々な実装があると思いますが，この節では螺旋本に載っている基本的な実装を紹介します．

本記事では触れることができませんでしたが，けんちょんさんの記事で紹介されている「めぐる式二分探索」もいずれ勉強してみたいと思います．二分探索を一般化し（キーが条件を満たすかどうか），螺旋本などで紹介されている基本的な二分探索の実装の分かりづらい部分を無くし，バグが少ない実装ができるそうです．

実装

アルゴリズム通りの基本的な実装です．

int BinarySearch(vector<int> vec, int key) {
    // key: 検索対象の値
    T left = 0, right = (T)vec.size(), mid;

    while (left < right) {
        mid = (left + right)/2;
        if (vec[mid] == key) {
            return mid;
        } else if (key < vec[mid]) {
            // 半分より下
            right = mid;
        } else if (vec[mid] < key) {
            // 半分より上
            // 必要ないが，分かりやすいように条件を記述してある．
            left = mid + 1;
        }
    }
    // key が配列の中に見つからない場合
    return -1;
}

要素を検索し一致したらインデックスを返す，という関数です．見つからなかった場合には -1 を返します．何を返すかは必要に応じてインデックスなのか bool なのか適宜決めてください．要素の格納には std::vector を用いており，この部分は配列で書き直すことができます．

実装上の注意点

実装の注意点は，

  • right の初期値が N-1 ではなく N であること
  • 条件分岐の 3 番目，「キーが半分より上」：vec[mid] < key のとき，left = mid + 1 だが，「半分より下」のときはright = mid で範囲の更新式に +1 が入らないこと

上記の注意点の部分を改変してみると，

  • 探索範囲の縮小が一定ステップ以上で起こらなくなり，探索が終了条件を満たせず無限ループに入ってしまう
  • 右端，左端の要素の探索が上手くいかない

など二分探索に不備が生じてしまいます．ぜひ手元で動かして確認してみてください．

C++ 標準ライブラリで二分探索

C++ では，標準ライブラリ algorithm をインクルードすることにより，便利な二分探索の関数が使えるようになります．二分探索関連で提供されている関数は主に以下の 3 つです．

  • std::binary_search()
  • std::lower_bound()
  • std::upper_bound()

二分探索を行うので，最悪計算量は O(\log N + 1) です．これらの使い方を解説していきます．

リファレンスへのリンクを下記に載せておきます．

std::binary_search()は，その名の通り二分探索による配列要素の検索を行います．使い方は

// first, last はソートした配列のイテレータ
// key は検索したい数値
std::binary_search(first, last, key)

のように，配列の先頭，末尾のイテレータと検索対象（数値，key）を渡します．戻り値は bool で，対象の数値があれば true，なければ false を返すという仕様です．最悪計算量は O(\log N + 1) です．

  • key の有無の判定を行う
  • key が見つかったとき，それがどの位置にあるかは分からない
  • 当然複数 key が存在した場合もどれに該当したのか分からない

という特徴があります．key と一致する要素に関する情報が「配列に存在するか否か」だけであり使い勝手が良いとは言えないです（その欠点を補う関数が後述する std::lower_bound()，std::upper_bound() になります）．

以下に std::binary_search() の簡単な使用例を示します．

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    vector<int> vec = {3, 1, 4, 6, 5};

    // ソート：二分探索はソートされた配列に対して適用する
    sort(vec.begin(), vec.end());  // 1, 3, 4, 5, 6

    // 0 - 10 の整数が vec に含まれているか検索
    for (int key = 0; key <= 10; ++key) {
        if (binary_search(vec.begin(), vec.end(), key)) {
            cout << key << ": " << "found" << std::endl;
        } else {
            cout << key << ": " << "not found" << std::endl;
        }
    }

    return 0;
}

出力

0: not found
1: found
2: not found
3: found
4: found
5: found
6: found
7: not found
8: not found
9: not found
10: not found

使い方はシンプルですね．

蛇足ですが，std::binary_search() を使用するときは必ずしもソートされている必要はありません．

検索対象より小さい数 < 検索対象 < 検索対象より大きい数

という並びが保証されていればソートをする必要はないという仕様になっています．上記の例の場合，本当はソートは必要ない（1 < 4 < 6 となっているため）のです．しかし，実際の問題ではそのような順序が保証されていることはないと思うので，とりあえずソート済みの配列に適用すると考えておいた方が良いと思います．

std::lower_bound() と std::upper_bound()

std::binary_search() は，ソートされた配列に対し key が存在するかしないかだけを返す，正直使い道が少ない不便な関数です．その欠点を補うための便利な二分探索の関数が std::lower_bound() と std::upper_bound() です．

どちらも std::binary_search() 同様ソートされた配列に対して二分探索を行う関数で，引数も同じ形を取ります．ソートされた配列の先頭，末尾のイテレータ，探索対象となる key の 3 つを渡して使います．

// first, last はソートした配列のイテレータ
// key は検索したい数値
std::binary_search(first, last, key)
std::lower_bound(first, last, key)
std::upper_bound(first, last, key)

用法

  • std::lower_bound() は，ソートされた配列に対して二分探索を行い，指定した key 以上の要素の内，一番左側の要素の位置（最小のインデックス）をイテレータで返します．

  • std::upper_bound() は，ソートされた配列に対して二分探索を行い，指定した key より大きい要素の内，一番左側の要素の位置（最小のインデックス）をイテレータで返します．

lower_bound() は「以上」，upper_bound() は「より大きい」*1となっています．どちらも要素自身ではなく，イテレータを返すことには注意が必要です．

std::lower_bound() は別の言い方をすると，指定された範囲（ソートされた範囲）における順序を崩さずに，key を挿入できる最初の位置を表すとも言えます．

基本

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    vector<int> vec = {1, 1, 2, 2, 4, 5, 5, 6, 8, 8, 8, 10, 15};
    vector<int>::iterator position;  // auto position; で良い．
    int idx_lower;

    position = lower_bound(vec.begin(), vec.end(), 3);  // 3 を二分探索
    idx_lower = distance(vec.begin(), position);
    cout << "vec[" << idx_lower << "] = " << *position << "\n";

    // cout << "vec[" << idx_lower << "] = " << vec[idx_lower] << "\n"; でも良い

    position = lower_bound(vec.begin(), vec.end(), 2);  // 3 を二分探索
    idx_lower = distance(vec.begin(), position);
    cout << "vec[" << idx_lower << "] = " << *position << "\n";

    return 0;
}

上記のコードの解釈をすると，

  • 3 以上の最初の要素の位置は，idx = 4 で，vec[idx] = 4 >= 3．
  • 2 以上の最初の要素の位置は，idx = 2 で，vec[idx] = 2 >= 2．

となります．これらの関数の戻り値はイテレータなので，わざわざポインタに代入しなくとも，直接 *upper_bound() のようにして要素にアクセスすることもできます．

upper_bound() と lower_bound() の比較

両者の比較をしてみます．

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    vector<int> vec = {1, 1, 2, 2, 4, 5, 5, 6, 8, 8, 8, 10, 15};
    vector<int>::iterator iter_lower, iter_upper;
    long idx_lower, idx_upper;
    for (int key = 0; key <= 16; ++key) {
        // lower_bound
        iter_lower = lower_bound(vec.begin(), vec.end(), key);
        idx_lower = distance(vec.begin(), iter_lower);

        // upper_bound
        iter_upper = upper_bound(vec.begin(), vec.end(), key);
        idx_upper = distance(vec.begin(), iter_upper);

        // output
        cout << "----- key = " << key << " -----" << "\n";
        cout << "lower_bound: vec[" << idx_lower << "] = " << *iter_lower << "\n";
        cout << "upper_bound: vec[" << idx_upper << "] = " << *iter_upper << "\n";
        cout << "\n";
    }
    return 0;
}

出力

全てを説明する必要はないと思いますが，key = 8 のときの結果を取り上げて説明します．結果は以下のようになりました．

# vec = {1, 1, 2, 2, 4, 5, 5, 6, 8, 8, 8, 10, 15}
    idx: 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12
----- key = 8 -----
lower_bound: vec[8] = 8
upper_bound: vec[11] = 10
  • lower_bound は「8 以上の要素の内，最小の要素は 8 でそのインデックスは 8」
  • upper_bound は「8 より大きいの要素の内，最小の要素は 10 でそのインデックスは 11」

以上のようにこれらの関数は，ソートされた配列を利用して二分探索を行い，特定の要素が存在する位置の特定を高速に行うことができます．

要素が見つからない場合

一つ注意点があります．配列の最大値よりも大きい key = 16 を探索したときの動作です．

std::lower_bound()，std::upper_bound() は，「key 以上（より大きい）の最初の要素を指すイテレータ」返しますが，条件を満たす要素が見つからなかった場合は，配列の末尾のイテレータを返します（std::vector なら vec.end()，配列なら A + N）．そのため，見つからなかった場合に戻り値のイテレータからインデックスを計算して使おうとすると配列の範囲外を指定して予期せぬ結果が起こる可能性があります．以下の例をご覧ください．

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    vector<int> vec = {1, 1, 2, 2, 4, 5, 5, 6, 8, 8, 8, 10, 15};

    long idx_last = distance(vec.begin(), lower_bound(vec.begin(), vec.end(), 16));

    cout << "vec[" << idx_last << "] = " <<  vec[idx_last] << "\n";  vec[13] = 0 が出力
    // cout << "vec[" << idx_last << "] = " <<  vec.at(idx_last) << "\n";  範囲外参照でエラー

    return 0;
}

key = 16 は vec 内に無いので std::lower_bound() は vec.end() を返します．vec.end() までの距離を std::distance によって算出すると idx_last = 13 となり，vec の範囲外となります．これにより idx_last を使って vec[idx_last] とアクセスすると 0 という予期せぬ値が返されます．範囲外アクセスでエラーを返す vec.at()を使うと当然実行時エラーとなります．

そのため，リファレンスなどでは，std::lower_bound()，std::upper_bound() の戻り値のイテレータを使うときに iter != vec.end() などで見つからなかった場合のイテレータを安全に処理している場合が多いです．ご注意ください．

std::lower_bound()，std::upper_bound() の活用

基本的な使い方を理解したところで，これらの活用例を見ていきます．

二分探索を行う

そもそも二分探索を行うための関数なので，当然二分探索で要素の検索を行うことができます．ここまでの説明が理解できていれば方法はわかるかと思います． 

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    vector<int> vec = {1, 1, 2, 2, 4, 5, 5, 6, 8, 8, 8, 10, 15};
    int key = 5;
    auto iter = lower_bound(vec.begin(), vec.end(), key);

    // 5 を検索
    if (iter != vec.end() && *iter == key) {
        cout << key << " is found at " << distance(vec.begin(), iter) << "\n";
    } else {
        cout << key << "is NOT found." << "\n";
    }

    return 0;
}

std::lower_bound() を使えば，「key 以上の最小の要素のイテレータ」を得ることができるので，そのイテレータから値を得て key と比較すればそれが存在するのかどうか簡単に二分探索を行うことができます．

注意点は，条件に値の比較だけでなく iter != vec.end() を含んでいることです．先述した通り要素が見つからなかった場合，自分で初期化等をしていない（値が未知）部分の配列末尾のイテレータが返ってきます（リファレンスに載ってませんでした）．そのため，末尾のイテレータに対して *iter でアクセスすると予測できない値が返ってくることがあります（そういう仕様なのかもしれませんが，私の環境では *vec.end() として 0 が返ってきます．

わざわざ配列末尾のイテレータにアクセスする処理を書くことがないためか，リファレンス等でも言及されていなかったので暇があったら()もう少し調べてみたいと思います）．

要素のインデックスを求める

本記事の中で度々使っていたのですが，std::distance() という関数を使うことで，二分探索で求めた「key と一致した要素」が配列のどこに位置しているのかを求めることができます（std::distance() の機能はイテレータ同士の距離を求めることです）．例で示しているので細かいソースコードは示しませんが，使い方だけ記述しておきます．

auto iter = lower_bound(vec.begin(), vec.end(), key);
unsigned int idx = distance(vec.begin(), iter);

配列先頭のイテレータと一致した要素のイテレータとの距離，つまり，インデックスを求めることができます．

ある数以上（より大きい）要素の数を数える

これは競プロの問題で使えるテクニックだと思います．

std::lower_bound()，std::upper_bound() を用いると数の境目が分かるため，それを利用して「配列の要素の内，key 未満（以下），より大きい（以上）要素数」を求めることができます．key に一致した要素のインデックスが分かるので当然と言えば当然の話ですが...

lower_bound()，upper_bound() の「以上」，「未満」に注意して， key が存在する位置を基準とした要素数の数え方の例を以下に示したので参考にしてください．

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    vector<int> vec = {1, 1, 2, 2, 4, 5, 5, 6, 8, 8, 8, 10, 15};
    size_t len = vec.size();
    int key = 6;

    // 二分探索
    auto iter_lower = lower_bound(vec.begin(), vec.end(), key);
    auto iter_upper = upper_bound(vec.begin(), vec.end(), key);

    // インデックスを取得
    long idx_lower = distance(vec.begin(), iter_lower);  // vec[idx_lower] = vec[7] = 6
    long idx_upper = distance(vec.begin(), iter_upper);  // vec[idx_upper] = vec[8] = 8

    // lower_bound で要素数を数える
    cout << "6 未満の要素数は " << idx_lower << "\n";  // 6 までの要素数 (idx_lower + 1) - 1
    cout << "6 以上の要素数は " << len - idx_lower << "\n";  // 6 までの要素数 (idx_lower + 1) - 1

    // lower_bound で要素数を数える
    cout << "6 以下の要素数は " << idx_upper << "\n";  // 6 までの要素数 (idx_lower + 1) - 1
    cout << "6 より大きい数の要素数は " << len - idx_upper << "\n";  // 6 までの要素数 (idx_lower + 1) - 1

    return 0;
}

二分探索で競プロの問題を解く

なんとなく活用の仕方が分かってきたと思うのでここからは，標準ライブラリを活用して競プロの問題を解いていきます．ネタバレがあるので先に問題を見たい方は以下に列挙したのでどうぞ．

AOJ: ALDS1_4_B - Binary Search

ただ二分探索をするだけの問題．基本を確認しましょう．線形探索をしようとすると制約時間内に間に合いません．

ABC061: C - Big Array

この問題は std::lower_bound() を使用しなくても解けるのですが，慣れてみたかったのであえて使って解いてみました．各数が何個あるかを集計して，その累積和が K 以上のときのインデックスが答えになります．

std::lower_bound() を使用する解法

std::lower_bound() を使用しない解法

正直こっちの方がとてもシンプルで見通しが良いです．公式解説もこちらで解いてます．

ABC084: C - Snuke Festival

らぴお（@Gxupu50ILbqonRQ）さんから教えていただいた問題．本記事で説明してきた std::lower_bound()，std::upper_bound() が大活躍の問題です．これらの関数を使うという事前知識があったからとはいえ，ちゃんと自力 AC 出来ました．下から順に積み上げようとすると計算量が O(N^{2} \log N) となってしまいますが，真ん中を軸として探索すると O(N \log N) で計算することができます．

難しめ？な問題

力尽きてしまったので問題の紹介だけします．解き次第追記しようと思います

補足：std::lower_bound() を使う上での注意点

Twitter で std::lower_bound() に関して調べていたらいくつか気になった言及があったので，それらを注意点としてまとめました．

algorithm の std::lower_bound() と std::set::lower_bound() の計算量の違い

algorithm の std::lower_bound() を std::set に対して使うと O(N) （線形時間になり，対数オーダーにならない）で，std::set のメンバ関数の std::set::lower_bound() 二分探索を行うと O(\log N) になります．つまり，std::set にはメンバ関数の方の lower_bound() で二分探索を行わないと意味がない（利点がない），ということです．

std::lower_bound() の計算時間が対数オーダーになるためには，渡すものがランダムアクセス可能なもの（配列，std::vector）である必要があり，ランダムアクセス可能でない構造の std::set （内部は平衡二分木）を渡すと対数オーダーにならなくてなってしまいます．

// O(N) になってしまう
std::lower_bound(set.begin(), set.end(), key)

// set に対して O(logN)
set.lower_bound(key)

これを知らずに，コンテスト中予期せぬところで計算量が O(N) になってしまい TLE をくらうと悲しいので、頭の片隅に置いときましょう．

メモ

本記事と関係ないことです．std::lower_bound() を調べている途中で偶然見つけて書くところが無かったのでメモさせてください．

std::count

std::count_if という関数を使うと，ラムダ式を使って条件を満たす要素の数を配列から簡単に計算することができます．

#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;

int main() {
    vector<int> v = { 1, 2, 3, 4, 4, 3, 7, 8, 9, 10 };

    // 3 と 5 の数を数える
    int key1 = 3;
    int key2 = 5;
    int num_key1 = count(v.begin(), v.end(), key1);
    int num_key2 = count(v.begin(), v.end(), key2);
    cout << "number: " << key1 << " count: " << num_key1 << '\n';
    cout << "number: " << key2 << " count: " << num_key2 << '\n';

    // ラムダ式を使って 3 で割り切れる数をカウント
    int num = count_if(v.begin(), v.end(), [](int i){return i % 3 == 0;});
    cout << "number divisible by three: " << num << '\n';
}

ラムダ式を知らなかったので少し勉強になりました．コンテスト本番だと間違えると嫌なので多分 for 文で書きます...

終わりに

最後まで読んでくださりありがとうございます．長すぎてグダッてしまったか感があります．記事中の間違いやおすすめの問題とかあれば Twitter で投げていただけると幸いです．

参考

*1:早く誰か未満と同じ感じの言葉を作ってくれ