PythonAlgorithm

Python bit全探索を簡単にやるコツ

公開日:2023-07-31 更新日:2023-07-31

久しぶりのブログ更新です。今日は競プロ関連の小ネタで、ちょっと簡単にbit全探索を実装する方法について書いていきます。

bit全探索とは

競技プログラミングで使うことのある全探索テクニックです。

基本的には、2つの選択肢がある問題の、すべての組み合わせを試すような探索方法です。

AtCoderなどのコンテストではC問題レベルで必要とされることがある、という印象です。

例題

bit全探索の基本的な考え方と、それを使用する際の簡単な例を示します。

例として、4つの数字 [1, 2, 3, 4] からいくつかを選んで、その和が特定の値 x になる組み合わせが存在するかどうかを考えます。

この場合、各数字を選ぶか選ばないかの2つの状態があるため、組み合わせの総数は2^4 = 16となります。

bit全探索を行うPythonの簡単なコードは以下のようになります。

nums = [1, 2, 3, 4]
x = 5  # 目的の合計

N = len(nums)
found = False

for i in range(2**N):
    total = 0
    for j in range(N):
        if (i >> j) & 1:  # j番目の要素を選ぶかどうか
            total += nums[j]
    if total == x:
        found = True
        break

print(found)

ポイントとしては 最初のループで、2を要素数だけ乗算します。この場合だと、2**Nですね。

そうして内側のループで要素数だけ繰り返し、i>>jして1が立っていたら、その要素を選ぶ、という具合です。

この式の >> は、ビット右シフト演算子です。

具体的には、i の2進数表現を j ビット分右にシフトします。

そして、その結果の最下位ビットが1か0かを確認するために & 1 というビット演算を行っています。

例えば、i = 5 および j = 2 の場合を考えます。 5の2進数表現は 101 です。

これを2ビット右にシフトすると、結果は 1 となります。

これを具体的に確認してみると、

iの2進数表現    : 101
0ビット右にシフト: 101
1ビット右にシフト: 010
2ビット右にシフト: 001
3ビット右にシフト: 000

末尾が1になっているのは0と2ビットのシフトの時です。

そのため、iが5の時は0番目と2番目の数字を選ぶ、ということになります。

結構難しい

書いてみたのはいいものの、覚えたての頃はこのi>>jのシフト部分が難しかったなーという記憶があります。

計算速度は落ちるんですが、シフト演算を使わずに考えられる方法があります。

Pythonではbin(n)とするとnを二進数表現の文字列に変換することができます。

bin(5)
>>>
'0b101'

先頭に0bがついてきます。

競プロなどの問題に使う時は2文字目以降をスライスして、先頭をゼロ埋めします。

bin(5)[2:].zfill(4)
>>>
'0101'

あとはこの文字列を走査して1が立っているところを考えれば良いです。

nums = [1, 2, 3, 4]
x = 5  # 目的の合計

N = len(nums)
found = False

for i in range(2 ** N):
    total = 0
    bin_str = bin(i)[2:].zfill(N)
    for j in range(N):
        if bin_str[j] == '1':
            total += nums[j]
    if total == x:
        found = True
        break
print(found)

実際の問題を解く

C - Skill Up

最大12までのN冊の参考書が売っています。高橋君の目標はM個あるアルゴリズム力をそれぞれX以上にすることです。

参考書の価格と、読むことで得られるアルゴリズム力が与えられるます。

目標を達成する最小の金額を出力する、という問題です。

この問題は上記のbit全探索を使って解くことができます。

具体的にはN冊の参考書を買う・買わないの組み合わせを全探索して、目標を達していた場合、金額を更新すればいいです。

n, m, x = map(int, input().split())

# 参考書の値段を記録
costs = []
# 参考書で得られるスキルを記録
books = []
for _ in range(n):
    c, *book = list(map(int, input().split()))
    costs.append(c)
    books.append(book)

ans = 10 ** 18
for i in range(2 ** n):
    # 0の時は一冊も買わないので、除外
    if i == 0:
        continue
    # 参考書を読んだ場合のスキル
    skills = [0] * m
    # 合計金額
    total_cost = 0
    # bit文字列
    bin_str = bin(i)[2:].zfill(n)
    for j in range(n):
        # j番目が1だったら買う
        if bin_str[j] == '1':
            # 本を取得
            book = books[j]
            # 価格を足す
            total_cost += costs[j]
            # スキルを足す
            for k in range(m):
                skills[k] += book[k]
    flg = True
    # 1個でも下回っていたらだめ
    for j in range(m):
        if skills[j] < x:
            flg = False
    if flg:
        # 答えを更新
        ans = min(ans, total_cost)
if ans == 10 ** 18:
    print(-1)
    exit()
print(ans)

おわりに

bit全探索系の問題はそもそもの制約が小さいことが多いので、多少計算量を犠牲にしてもほぼ間に合います。

bin(n)で二進数文字列にする方法は覚えておいて損はないでしょう。ではでは。

Twitter Share