めぐる式二分探索をPythonで実装する。シンプルな全探索との比較あり。

例えばAtCoderのこの問題　(abc146のC)

線形探索で解こうとすると、探索範囲が$ 1≤A≤10^{9} $ となるので時間内に計算が終わりません。
（※$ 10^{8}$から$10^{9} $程度でTLEになります。）

試しに実装してみましょう。

A,B,X = map(int,input().split())

max_int = 0
for i in range(1,10**9+1):
    if (A * i) + ( B * (len(str(i))) ) <= X:
        max_int = i
    else:
        break

print(max_int)

愚直に１から調べ上げる実装ですが、これを提出すると想定通り、TLEになってしまうものが大半です。このアルゴリズムをみてみると、1から順番に1ずつ大きくしていって条件を満たすかチェックしてます。

ということは$1,2,3,4,5,6,, ,,10^{9}$というソート済みの配列に対して探索をしていることと同じです。このようなソート済みの時に使えるアルゴリズムといえば、、、。そう、二分探索ですよね。

ではアルゴリズムの教科書に書いてあるような方法で実装してみましょう。

A,B,X = map(int,input().split())

#探索したい整数の範囲を定義
left = 1
right = 10 ** 9

#満たすべき条件
def validate(arg):
    if (A * arg) + ( B * (len(str(arg))) ) <= X:
        return True

#条件を満たす最大整数を探す二分探索
max_int = 0
while left <= right:
    mid = (left+right) // 2
    if validate(mid):
        max_int = mid
        left = mid + 1
    else:
        right = mid -1

print(max_int)

通常の2分探索であれば、条件を満たす値が見つかった時点でWhile文を抜けますが、この問題は条件を満たす最大の整数をみつける問題なので、そのまま続けます。ただし条件を満たす場合だけ、max_intを更新するのがポイントです。このコードのmidを答えに使うとバグります。

これでも問題ないのですが、めぐる式で指摘されているように、mid + 1 や mid -1 の部分で気をつけないと最後の境界判定のところでバグりやすいです。またWhileの条件判定も、<なのか<=なのか、間違えやすいポイントでもあります。

そこで、めぐる式を参考にして書き換えるとこうなります。

※めぐる式を参考にすると right-left > 1 となるまで条件を満たす最大のleftを探せばOKということになります。詳しくはこちら

※※めぐる式本家こちらです

A,B,X = map(int,input().split())

#探索したい整数の範囲を定義
left = 0
right = 10 ** 9 + 1

#満たすべき条件
def validate(arg):
    if (A * arg) + ( B * (len(str(arg))) ) <= X:
        return True

#条件を満たす最大整数を探す二分探索
max_int = 0
while abs(right - left) > 1:
    mid = (left+right) // 2
    if validate(mid):
        max_int = mid
        left = mid
    else:
        right = mid

print(max_int)

めぐる式の場合、探索したい範囲の上限、下限より１つずつ大きくするのがポイントです。
このコードの場合、left = 0　　right = 10 ** 9 + 1　としてやります。

たいした変化ではないように見えますが、+1 -1がなくなるのと、>= なのか＞なのかで悩まなくてもすみますね。

abs(right - left)の差が１になったら終了です。