二分探索サンプルコード集(コピペ用)

二分探索は、感覚的なわかりやすさに反してバグが入りやすいことで有名なアルゴリズムです。20 の教科書のうち 15 でバグっていたという報告もあるそうです。

実際、自分も書くたびにバグに苦しんできました。変な値を返すだけでなく、out of bounds アクセスや無限ループもよく起きます。一旦動いたと思っても、後になってバグが発症することも多く、たちが悪いです。

そこで、きちんとテストした二分探索のサンプルコードを自分のコピペ用に作ってみました。

動作仕様 (境界探索版)

ソートした配列 a に対して、「値が c 以上になる範囲のうちの一番左のインデックス」を返す関数 bsearch_min を書きます。

a = [0, 1, 1, 1, 2, 2, 2, 3]
p bsearch_min(a, 2) #=> 4

値が c 以上になる値がない場合は a.size を返します。空配列の場合は 0 を返します。

p bsearch_min(a, 4)  #=> 8 (a.size と同じ)
p bsearch_min([], 0) #=> 0 (a.size と同じ)

よって、bsearch_min の返り値は 0 以上 a.size 以下です。

なお、「値が c 未満になる範囲のうちの一番右のインデックス」が欲しくなることもありますが、この場合は bsearch_min の返り値から 1 を引けば OK です(その場合の返り値は必然的に -1 以上 a.size 未満になります)。どちらかというと、「左端」と「右端」のどちらが欲しいのかを混乱しないことが重要です。

区間を使う bsearch_min

def bsearch_min(a, c)
  l, r = 0, a.size - 1
  while l <= r
    m = l + (r - l) / 2
    if a[m] < c
      l = m + 1
    else
      r = m - 1
    end
  end
  l
end
  • l と r の扱いが対称的で美しい。
  • 終了時、l - 1 == r となっている。よって、「左端」ではなく「右端」が欲しくなったときは全く同じコードで r を使えばよい。
  • 除算の丸めに対してロバスト。つまり m = l + (r - l + 1) / 2 としても正常に動作する。
  • l = m + 1 と r = m - 1 から、イテレーションのたびに範囲が狭くなることが見た目に明らか。二分探索は無限ループバグを入れやすいので重要。
  • m が求めるインデックスだった場合にも r = m - 1 を実行して大丈夫なのか不安になる。つまり、求めているインデックスを i としたとき、i が (l..r) の範囲外になることがある。たとえば bsearch_min([1, 2, 3], 2) を実行すると、求めるインデックスは 1 なのに、(l..r) は (0..2) → (0..0) になるので、1 は範囲外になる。その後で l が r を追い越して (l..r) = (1..0) になり、l を返すので、バグではない。

半開区間を使う bsearch_min

def bsearch_min(a, c)
  l, r = 0, a.size # 変更点 (1)
  while l < r # 変更点 (2)
    m = l + (r - l) / 2
    if a[m] < c
      l = m + 1
    else
      r = m # 変更点 (3)
    end
  end
  l
end
  • 違いは、(1) r = a.size - 1 から r = a.size になった、(2) ループ条件が l <= r から l < r になった、(3) r = m - 1 から r = m になった、の 3 点。
  • l と r の扱いが対称的でない。半開区間だから当然だけど。
  • 終了時に l == r になっている。混乱しないという意味では利点? ただし、「右端」が欲しくなったら l - 1 を計算する必要がある。
  • 除算の丸めは floor でないといけない。(半開区間なので、すでに r に 1 足されてるようなもの)
  • r = m なので、無限ループバグが怖くなる。ループ中は l < r なので、除算が floor なら必ず m < r になる、よって範囲は確実に狭まる。というように考えれば大丈夫だとわかるんだけど、二分探索は脳力を使うところが多いので疲れる。
  • 求めているインデックスは必ず (l..r) の範囲にある。ただし、半開区間を閉区間として解釈すれば、という話なので余計にややこしい気もする。

動作仕様 (有無判定版)

ソートした配列 a に対して、「値 c を含むかどうか、含む場合はそのインデックス」を返す関数 bsearch_any を書きます。

a = [0, 1, 1, 1, 2, 2, 2, 3]
p bsearch_any(a,  2) #=> 4 or 5 or 6
p bsearch_any(a, -1) #=> nil
p bsearch_any(a,  4) #=> nil

該当する値が複数ある場合は、どれかのインデックスを返します。

区間を使う bsearch_any

def bsearch_any(a, c)
  l, r = 0, a.size - 1
  while l <= r
    m = l + (r - l) / 2
    return m if a[m] == c # c と等しい要素を見つけたら return
    if a[m] < c
      l = m + 1
    else
      r = m - 1
    end
  end
  nil # 見つからなかった
end

境界探索版との違いは、return 文を入れたことと、最後に nil を返すところだけです。

半開区間を使う bsearch_any

def bsearch_min(a, c)
  l, r = 0, a.size
  while l < r
    m = l + (r - l) / 2
    return m if a[m] == c # c と等しい要素を見つけたら return
    if a[m] < c
      l = m + 1
    else
      r = m
    end
  end
  nil # 見つからなかった
end

こちらも同様。

講評

「閉区間より半開区間を使う方がプログラムが綺麗になることが多い」という経験則を持ってる人は多いと思いますが、こと二分探索に関しては閉区間に分があるように思いました。

二分探索というと、境界探索より有無判定を指すことが多いようです。が、個人的には二分探索を使いたい場合には境界探索をしたいことが多い気がする。特に、「複数ある場合はどれでもいい」というケースは出会った記憶がない(重複がないとわかってる場合に使う?)。あと、境界探索版を有無判定版にするのはミスしにくいんだけど、逆は脳力を要する(a[m] == c の場合に l と r のどちらを変えるべきか、最終的な返り値は l と r のどちらか)ので、やはり境界探索版をベースだと思っておいたほうが良いと思う。

なお、Ruby なら組み込みの Array#bsearch や Range#bsearch を使うべきです。この記事はあくまで、何かやむを得ない事情で自力実装が必要になったとき用の備忘録です。

以下、テストに使ったコード。隣同士の差が 0 から 2 になる配列を長さ 0 から 8 まで網羅的に試しています。

def bsearch_min_closed(a, c)
  l, r = 0, a.size - 1
  while l <= r
    m = l + (r - l) / 2
    if a[m] < c
      l = m + 1
    else
      r = m - 1
    end
  end
  l
end

def bsearch_min_half_open(a, c)
  l, r = 0, a.size
  while l < r
    m = l + (r - l) / 2
    if a[m] < c
      l = m + 1
    else
      r = m
    end
  end
  l
end

def bsearch_any_closed(a, c)
  l, r = 0, a.size - 1
  while l <= r
    m = l + (r - l) / 2
    return m if a[m] == c
    if a[m] < c
      l = m + 1
    else
      r = m - 1
    end
  end
  nil
end

def bsearch_any_half_open(a, c)
  l, r = 0, a.size
  while l < r
    m = l + (r - l) / 2
    return m if a[m] == c
    if a[m] < c
      l = m + 1
    else
      r = m
    end
  end
  nil
end

0.upto(8) do |nn|
  [0, 1, 2].repeated_permutation(nn) do |diff|
    next if diff.last != 0
    a = []
    n = 0
    diff.each do |d|
      a << n
      n += d
    end
    p a
    -2.upto((a.max || 0) + 2) do |c|
      answer = a.find_index {|v| v >= c } || a.size
      raise if bsearch_min_closed(a, c) != answer
      raise if bsearch_min_half_open(a, c) != answer
      if a.include?(c)
        raise if a[bsearch_any_closed(a, c)] != c
        raise if a[bsearch_any_half_open(a, c)] != c
      else
        raise if bsearch_any_closed(a, c) != nil
        raise if bsearch_any_half_open(a, c) != nil
      end
    end
  end
end