【AtCoder】Pythonで二分探索

はじめに

Pythonで二分探索を実装するためのテンプレとAtCoderの問題をまとめました。

Pythonで二分探索

二分探索とは、検索する問題を半分に分割しながら解を探しだすアルゴリズムです。

Pythonでソートされたリストに対して、ソートされた状態を保ちながら挿入・挿入できる場所を求める場合はbisectというライブラリが便利です。下記の記事で解説しています。

【AtCoder】bisectでリストを二分探索する
はじめにbisectの使い方とbisectを使うAtCoderの問題をまとめました。bisectの使い方bisectは、ソートされたリストにソートされた状態を保ちながら挿入、挿入する場所を求めることができるライブラリです。二分探索を...

本記事で紹介する二分探索はbisectを利用せず、ある範囲から解を探すための実装になります。

二分探索のテンプレ

bisectを利用しない二分探索のテンプレは以下になります。

def is_ok(mid):
    """二分探索中の判定

    Parameters
    ----------
    mid : int

    Returns
    -------
    resutl : bool
    """
    if mid > 0:
        return True
    else:
        return False

def binary_search(ok, ng):
    """二分探索

    Parameters
    ----------
    ok : int
    ng : int

    Returns
    -------
    ok : int
    """
    while ng - ok > 1:
        mid = (ok + ng) // 2
        if is_ok(mid):
            ok = mid
        else:
            ng = mid
    return ok

binary_search関数で検索する範囲を半分にしながら解を探します。

is_ok関数の判定基準を修正することでいろんなパターンに対応できるようにしています。

AtCoderのサンプル問題

ZONeエナジー プログラミングコンテスト “HELLO SPACE” C問題

C - MAD TEAM
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

N人から3人選んでチームの総合力を最大化する問題です。

is_okの中で探索中の値以上になるか判定しています。

n = int(input())
ability = [list(map(int, input().split())) for _ in range(n)]

def is_ok(mid):
    """二分探索中の判定

    Parameters
    ----------
    mid : int

    Returns
    -------
    resutl : bool
    """
    s = set()
    for a in ability:
        s.add(sum(1 << i for i in range(5) if a[i] >= mid))
    for x in s:
        for y in s:
            for z in s:
                if x | y | z == 31:
                    return True
    return False

def binary_search(ok, ng):
    """二分探索

    Parameters
    ----------
    ok : int
    ng : int

    Returns
    -------
    ok : int
    """
    while ng - ok > 1:
        mid = (ok + ng) // 2
        if is_ok(mid):
            ok = mid
        else:
            ng = mid
    return ok

ans = binary_search(0, 10**9+1)
print(ans)

ABC146 C問題

C - Buy an Integer
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

所持金X円の時に、買うことができる最も大きい整数を求める問題です。

シンプルな二分探索の問題です。

a, b, x = map(int, input().split())

def is_ok(mid):
    """二分探索中の判定

    Parameters
    ----------
    mid : int

    Returns
    -------
    resutl : bool
    """
    if (a * mid) + (b * len(str(mid))) <= x:
        return True
    else:
        return False

def binary_search(ok, ng):
    """二分探索

    Parameters
    ----------
    ok : int
    ng : int

    Returns
    -------
    ok : int
    """
    while ng - ok > 1:
        mid = (ok + ng) // 2
        if is_ok(mid):
            ok = mid
        else:
            ng = mid
    return ok

ans = binary_search(0, 10**9+1)
print(ans)

ABC192 D問題

D - Base n
AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

Xをn進数とみなして得られる値で、M以下になるものは何種類あるか求める問題です。

n進数から10進数に変換しながらis_okでM以下か判定します。

x = list(input())
m = int(input())
X = int(''.join(x))
max_x = int(max(x))

def change_nto10(num, b):
    n = 0
    numlist = list(str(num))
    while(numlist):
        n *= b
        n += int(numlist.pop(0))
    return n

def is_ok(mid):
    """二分探索中の判定

    Parameters
    ----------
    mid : int

    Returns
    -------
    resutl : bool
    """
    t = int(change_nto10(X, mid))
    if t <= m:
        return True
    else:
        return False

def binary_search(ok, ng):
    """二分探索

    Parameters
    ----------
    ok : int
    ng : int

    Returns
    -------
    ok : int
    """
    while ng - ok > 1:
        mid = (ok + ng) // 2
        if is_ok(mid):
            ok = mid
        else:
            ng = mid
    return ok

if len(x) == 1:
    if X > m:
        ans = 0
    else:
        ans = 1
else:
    ans = binary_search(max_x, 10**18+1) - max_x

print(ans)

参考

タイトルとURLをコピーしました