二分探索の基本的な問題を解く

f:id:hrksb5029:20210306001143p:plain
二分探索の基本的な問題を解いて練習します。
自分なりに思ったポイントも説明します。

問題はAIZU ONLINE JUDGEのIntroduction to Algorithms and Data StructuresのBinary Searchです。
https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_4_B&lang=ja

ざっくりと問題をいうと、2つの数列を読み込んで、2つの数列で共通する数字の個数を答える問題です。

Pythonだとsetがあって、setの積をとって解きたくなってしまいますが、今回は練習なので二分探索で解きます。

二分探索をざっくり解説すると、ソートされた数列(ここでは左が小さくて右の方が大きいと仮定)の真ん中の値と探している数字を比較します。
真ん中の値の方が大きければ数列の大きい方(右側)を探して、小さければ数列の小さい方(左側)を次に探します。
今度はその中で真ん中の値と探している数字を比較します。
このように探索する範囲を半分ずつに絞り込みながら探すアルゴリズムです。

半分ずつなので、数列の長さが2のとき数字を見つけるまでに真ん中の値と比較する回数は1回、
数列の長さが4のとき比較する回数は2回、
数列の長さが8のとき比較する回数は3回
となるので、計算量は O(log_2 N)です

実装するときのポイントは3つあります。

探す範囲のインデックスを記録する

探す範囲の下限と上限のインデックスを記録します。
探す範囲の下限の初期値は0で、上限の初期値は探す対象の数列の[長さ-1](インデックスなので範囲は0〜N-1)です。
そして真ん中の数字となるインデックスは(high + low) // 2になります。

while文の条件

while文の条件は探している値が数列の中に見つからなかったときにループから抜けられるものにします。
下限のインデックス値が上限のインデックス値を逆転したら終わりです。
while low <= high:
となります。
他にもいろいろな書き方はあると思います。

探す範囲の更新

中央の値と比較した結果に応じて、上限のインデックス値と下限のインデックス値を更新します。
コードで表現するとこんな感じです。

if S[mid] == t:
    count += 1
    break
elif S[mid] > t:
    high = mid - 1
else:
    low = mid + 1

最後にまとめたコードはこうなります。

N = int(input())
S = list(map(int, input().split()))
S.sort()
Q = int(input())
T = list(map(int, input().split()))
count = 0
for t in T:
    high = N - 1
    low = 0
    while low <= high:
        mid = (high + low) // 2
        if S[mid] == t:
            count += 1
            break
        elif S[mid] > t:
            high = mid - 1
        else:
            low = mid + 1
print(count)

個人的な好みとしてはelseやelifを書くのはあまり好きではなくて、breakのあとはif文にして、elifの中にcontinueを書いて、elseは削除します。

N = int(input())
S = list(map(int, input().split()))
S.sort()
Q = int(input())
T = list(map(int, input().split()))
count = 0
for t in T:
    high = N - 1
    low = 0
    while low <= high:
        mid = (high + low) // 2
        if S[mid] == t:
            count += 1
            break
        if S[mid] > t:
            high = mid - 1
            continue
        low = mid + 1
print(count)