🟫

AtCoder灰・茶・緑色の方必見!二分探索を絶対にバグらせないで書く方法

2019/12/23に公開

この記事はCompetitive Programming (1) Advent Calendar 2019 23日目の記事です。

旅行プラットフォーム事業部の大沢です。

競技プログラミングを2年前に始めて以来、週末のAtCoderコンテストにはほとんど欠かさず出ています。
私は昨年末に青色コーダーになり、実力をどうにかキープしています。まだ時間はかかってでも強くなりたい気持ちがあります。

この記事の気持ち

二分探索についての教材は世の中に多くあり、良質な記事も多い反面、「半開区間」などの考え方が難しく混乱するという意見も耳にしています。また、実際に書いてみると意外とバグりやすいことでも有名で、私もよくハマってしまうことがありました。「半開区間」という言葉を使わず、私なりにわかりやすいと思う理解と、バグりにくい書き方を記事にしてみました。

メインターゲットの読者は、以下のいずれかを想定しています。

  • 二分探索って何だろう?という方
  • 二分探索の概要をなんとなく知っている方
  • 二分探索を学んだことがあるが、理解がちょっと怪しい方
  • 二分探索を実際に書いたことがあるけれど、よく細部をバグらせてしまう方
  • 二分探索をあらためて直感的に理解したい方

特に、細部をバグらせないような、直感的でおすすめな理解の仕方を紹介したいと思います。

二分探索とは

二分探索は一言で言うと「境目を見つける」アルゴリズムです。

探索範囲の1か所に境目があって、「境目の左側が全てある条件を満たし、右側が全てその条件を満たさない」ことがわかっているときに、その境目を高速に見つけることができます。
もちろん条件を満たす満たさないは左右逆でも使えます。

境目というのは何でもよいです。読みかけの本のここまで読んだ/読んでないのページの境目とか、背の順に並んだ児童の中で身長が100cm未満/以上の境目とか、納期に間に合う/間に合わないのタスク量の境目とか、全社員を満足させるために足りる/足りないのピノの箱数の境目、とか・・・。

とにかく「1か所の境目」の両側で判定結果が二分されていることが重要です。判定結果が「1か所の境目」で二分されない条件では基本使えないと思ってください(使えないことは無いですがこの記事では扱いません)。

実装は後ほど解説しますが、まずは長さ10の配列を使って、二分探索の動き方を見ていきましょう!

以下の例では左側が条件を満たす側だとします(逆の場合は後述します)。

図1

図1で、黒い枠線は要素数10の配列だとします。上の緑文字がindexです。
青い領域は条件を満たすことが確定した領域、赤い領域は条件を満たさないことが確定した領域です。

  1. ok,ngの2つの変数を用意し、探索範囲の外になるような値を設定する※1
    • ok=-1, ng=10と置く
  2. okngの平均を求める。4となる
    • is_ok(4) == Trueとなり条件を満たすのでok = 4とする※2
  3. okngの平均を求める。7となる
    • is_ok(7) == Falseとなり条件を満たさないのでng = 7とする
  4. okngの平均を求める。5となる
    • is_ok(5) == Trueとなり条件を満たすのでok = 5とする
  5. okngの平均を求める。6となる
    • is_ok(6) == Falseとなり条件を満たさないのでng = 6とする okngの差が1になったので処理を終了する

※1 ok,ngの初期値について 探索範囲の思いっきり外側でもよいし、探索範囲の内側でもよい。 大事なのはokは確実に条件を満たすゾーン、ngは確実に条件を満たさないゾーンに含まれていること。これが間違っていると正常に動作しません。
※2 判定関数 is_ok(i)はindexがiのときに条件を満たすならTrue、そうでなければ Falseを返します。配列外のi が引数で来たときにも、満たす側ならTrue、満たさない側ならFalseと返すものとします。 今回の例では単純に、

def is_ok(i):
   return i <= 5

のような実装がされていると思ってください(ここでは配列の中身すら無視されていますが・・・)。

理解のポイントとしては、

  • okは常に条件を満たすことが確定したゾーンの一番右側にいる
  • ngは常に条件を満たさないことが確定したゾーンの一番左側にいる
  • 最終的にokngは密着する(差が1になる)ことがわかれば完璧です!

この挙動をするコードをPythonで書くと次のようになります。

ok = -1
ng = 10
while ng-ok > 1:
    mid = (ok+ng) // 2 # 平均(小数切り捨て)
    if is_ok(mid):
        ok = mid
    else:
    	ng = mid
print(ok,ng) # "5 6" が出力される

それでは、5で最終結果として得られたok,ngの値は何を示すでしょうか?

  • ok・・・条件を満たすなかで最大のindex
  • ng・・・条件を満たさないなかで最小のindex

この理解でほとんど問題ありません(注意すべき点は後述します)。
実際okとして得られた5は、is_ok(i)を満たすiのうち最大の整数です。

左側がngの場合

また、左側がngとして実装した場合、コードは例えばこのようになります。

def is_ok(i):
    return i > 5 #大きい側がTrue
 
ok,ng = 10,-1 # さっきと逆なので注意
while ok-ng > 1: # さっきと逆なので注意。abs(ok-ng)のように汎用的に書く流派もある
    mid = (ok+ng) // 2 # 平均(小数切り捨て)
    if is_ok(mid):
        ok = mid
    else:
    	ng = mid
print(ok,ng) # "6 5" が出力される

そして、

  • ng・・・条件を満たさないなかで最大のindex
  • ok・・・条件を満たすなかで最小のindex

となります。

ここでよく混乱しがちなのが、二分探索で得られた2つのポインタのうち、最終的にどちらを使えばよいのか?という問題です。

大体の場合、okを使えばOK

ところで、このような問題文をよく見ませんか?

  • 条件を満たすなかで最大の〇〇を求めよ
  • 条件を満たすなかで最小の〇〇を求めよ

このような問題文が出てきたときは、二分探索を使えるケースが少なくないです。
そして、

  • 条件を満たすなかで最大の〇〇 → okを左側として実装し、最終的にokを使う
  • 条件を満たすなかで最小の〇〇 → okを右側として実装し、最終的にokを使う

なんと、okとして得られた値をそのまま使えばよいのです!
「条件を満たす側」「条件を満たさない側」と分けてきたのは、このためです。

一般的な二分探索では2つのポインタをhigh, lowみたいな名前で管理することが多いと思うのですが、ok, ngとすることで、何を扱っているのかがわかりやすくなり、何かと嬉しいことが多いです。

このok,ngで管理する方式は私が考えたのではなく、いわゆる「めぐる式二分探索」として知られています。

配列外参照には注意

  • 配列の要素すべてが条件を満たさない場合、ok = -1となり、okが配列外を指します。
  • 同様に、要素すべてが条件を満たす場合、ng = 10となり、ngが配列外を指します。

左をngとするケースではこれの逆で、ng = -1ok = 10の状態が生じます。

私がよくやるのは、下のような関数を作って判定します。

def is_ok(i):
	if i < 0: return True
	if i >= N: return False
	return 有効なiに対する判定

引数のiが配列外など、有効な範囲にないときの処理を忘れないようにしましょう。
また、向きにも注意で、okを返す側の異常値のときに Truengを返す側の異常値のときにFalseを返してください。

「個数を求める」場合などもちょっと注意

  • 条件を満たすものがいくつあるか?

などという問題に対してはちょっとだけ注意が必要です。

結論を書くと、

  • 条件を満たす側が左側なら → ngが答え
  • 条件を満たす側が右側なら → N - okが答え

になります(が、これは覚える必要はありません)。

あくまでokが指すのはギリギリ条件を満たすボーダーの入力(ここではindex)です。
indexから個数を求める必要があるということは頭の片隅に置いてください(これが頭から抜けていると、二分探索の実装がバグっているのか?と錯覚して焦りがちです)。

なぜこのように求まるかはぜひ考えてみてください。次節のような図を描けば、感覚的にも理解できるかと思います。

最終的なok,ngについて視覚的な理解

ところで、要素数10の配列を二分探索した場合、得られる結果は11通りあります。
要素数NならN+1通りです。

図2

上の図に書かれた 010の青い数字は、要素の境目に番号を振ったものです。
このように0始まりで番号を振った場合、二分探索の結果のうち、ok,ngの大きい方に一致します。

たとえば

  • ok=5, ng=6ならば6という境目
  • ok=3, ng=2ならば3という境目

が求められたことを示します。

これがイメージできていれば、最終的に二分探索によって何が得られているのかが確実に理解できているはずです!図で捉えればもう怖いものはないですね!

二分探索の強力さ

探索範囲がN要素の場合、log_2(N)回程度の比較回数でよいです。

例であげたような10要素ぐらいの場合では効果は薄いのですが(むしろ二分探索を使わない方が実装が単純な分よい)、探索範囲の規模が大きくなるほど効果が強いです!

例えば

  • 100000 (=10^5)要素 → 17回程度
  • 5000兆 (=5*10^15)要素 → 53回程度

の比較回数で求まりかなり強力です!

図1(再掲)

図1を改めて見ていただくと、未確定のゾーン(白い部分)が1回の処理でおよそ半分にしぼり込まれていく様子がわかると思います。

先ほど5000兆要素と書きましたが、実際にこれだけの長さの配列がメモリに乗ることは現実にはないと思います。
実は、探索対象は配列でなくてもよいのです。

is_ok(x)関数の結果の、True / Falseが切り替わる境目が1か所以下(=単調性がある)ならば、二分探索が使えます。

探索対象が浮動小数点数をとる場合

お気づきの方もいるかと思いますが、探索対象が配列でなくてもよいということは、is_ok(x)の引数x整数以外を取ってもよいということです。
実際、境目となる浮動小数点数の値を二分探索で調べたいケースもあります。

while ng-ok > 1:でループさせるような先ほどの実装では期待した動きになりません。

このようなときは、何も考えずに100回程度ループするのが定石となっているようです。
(一般的なdouble型の精度より遥かに余裕がある回数なので、お好みで調整してください。)

for i in range(100):
	mid = (ok+ng) / 2 # 平均(浮動小数点数)
    if is_ok(mid):
        ok = mid
    else:
    	ng = mid

このケースではokngは同じ値に収束してくるので、どちらを使う?のように考える必要はないですね。

bisectモジュールの使えるところと使えないところ

※これはPythonista向けの話題です。

Pythonには、ソートされた配列の中に、ある値が入るべき境界を二分探索で見つけたいとき、標準モジュールにbisect.bisect_left()bisect.bisect_right()といった関数があります(C++ だとstd::lower_bound()std::upper_bound()に相当する気がします)。

非常に便利ですので、これらについて最後に簡単に紹介します。

まず使用例です。

from bisect import bisect_left, bisect_right
 
arr = [1,3,5,5,5,6,7] # 昇順にソートされている必要がある
l = bisect_left(arr, 5)  # 5が入るべき境目のうち最も左側の境目を返す
r = bisect_right(arr, 5) # 5が入るべき境目のうち最も右側の境目を返す
print(l,r) # "2 5" が出力される

このように、昇順ソートされた配列の境目を見つけるタイプの問題であれば、bisectモジュールを呼び出せば自分で実装する必要はありません。私も時々これのお世話になっています。

ちなみにこれの返り値の正体は、図2で書かれている青い数字に相当するものが返ってきます。
これがわかっているだけでもライバルに差がつきます!

また、配列の中身が数値でなくてもbisectは利用できます。文字列やタプルの場合も辞書順比較をしてくれます。Comparableな要素の配列で、正しく昇順ソートされていればよしなにやってくれます。

では逆に使えないケースはどんなときかというと、探索範囲が配列ではないときです。
関数に入力されるxのうち、条件を満たす/満たさないxのボーダーを見つけたい場合は、自前実装するしかありません。

xについて条件を満たすかどうかの判定が配列の値以外によるのであればこれに該当します。
(配列でなく関数の結果を探索する場合、二分法と呼ぶのが正しい気がしていますが、競プロの文脈では区別されないことが多いです。この記事でも「二分探索」に統一して呼んでいます。)

降順ソートされた配列に対しては、自前実装で対応してもよいのですが、私は反転した配列にbisectを使い、それの結果を反転することが多いです。
配列の反転にはO(N)かかりますが、入力の時点でO(N)かかっているはずなので問題になることはおそらくありません。

まとめ

  • okは条件を満たすことが確定したゾーン、ngは満たさないことが確定したゾーン
    • どちら側が条件を満たすのか、には要注意
  • okとngの2変数で未確定のゾーンをはさみながら絞り込んでいく
  • 最終的には ok を使おう
    • ただしそのまま使えないケースもあるので、よく見極めて
  • okngの境目を求めているのだということを理解しよう
  • 図でイメージできれば、もう間違えない!

水色になってぜひフォルシアへ!

フォルシアでは2021年度新卒採用を行っています。

近年AtCoderJobsからの応募・入社が増えてきています。

強い技術を持ちながらビジネスに活かしたいWebエンジニアの方、水色以上になってAtCoderJobs経由で応募いただくと書類選考が免除されます!
仕様を実装に落とし込むのが早く、計算量の感覚も身についていると、とても素敵です!

競プロerの皆さん、ぜひ一緒に働きましょう!

FORCIA Tech Blog

Discussion