Pythonで解く【初中級者が解くべき過去問精選 100 問】(010 - 014 全探索：ビット全探索)

1. 目的

初中級者が解くべき過去問精選 100 問をPythonで解きます。
すべて解き終わるころに水色になっていることが目標です。

本記事は「010 - 014 全探索：ビット全探索」です。

2. 総括

今まではbit全探索の際はitertools.productで(0, 1)を生成していましたが、後々のためにbitシフトでも解いてみました。

3. 本編

010 - 014 全探索：ビット全探索

010. ALDS_5_A - 総当たり

回答

n = int(input())
A = list(map(int, input().split()))
q = int(input())
M = list(map(int, input().split()))

check_list = [0] * (2000 * 20 + 1)
for i in range(2**n):
    total = 0
    for j in range(n):
        if (i >> j) & 1:
            total += A[j]
    check_list[total] = 1

for m in M:
    if check_list[m]:
        print('yes')
    else:
        print('no')

リストAのすべての要素について「使う」「使わない」を判定し、その合計がリストMに含まれるか否かを判定します。
最初は下記のようなコードとしましたが3重のforループとなり間に合いません。
算定したtotalを毎回mと比較すると速度が遅くなるため、算定したtotalを添え字とするcheck_listにフラグをたて、そのあとに改めて添え字mをチェックする必要があります。

# これだと間に合わない
n = int(input())
A = list(map(int, input().split()))
q = int(input())
M = list(map(int, input().split()))
for m in M:
    ans = 'no'
    for i in range(2**n):
        total = 0
        for j in range(n):
            if (i >> j) & 1:
                total += A[j]
        if total == m:
            ans = 'yes'
            break
    print(ans)

---

011. AtCoder Beginner Contest 128 C - Switches

回答

N, M = map(int, input().split()) # Nはスイッチの数、Mは電球の数
lights = [[0] * N for _ in range(M)]
for i in range(M): 
    temp = list(map(int, input().split())) # 0番目はスイッチの個数、1番目以降はスイッチを示す
    k = temp[0]
    switches = temp[1:]
    for j in range(k):
        lights[i][switches[j]-1] = 1
P = list(map(int, input().split())) # 個数を2で割ったあまりが要素と等しい場合に点灯する

answer_count = 0
for i in range(2**N): # bit探索
    flag = True
    for k in range(M): # すべての電球について調べる
        count = 0
        for j in range(N): # bitで1フラグがある番号について行う
            if (i >> j) & 1:
                count += lights[k][j]
        if count % 2 != P[k]: # 1個の電球でもpと一致しなかった場合はflagをFalseにしてbreak
            flag = False
            break
    if flag: # すべてクリアしてflagがTrueの場合にcountを増やす
        answer_count += 1

print(answer_count)

問題文を理解するのが若干難しいです。
入力の与えられ方が扱いずらいので、入力の受け取り方を工夫し、縦軸を電球、横軸をスイッチの2次元配列を作成しました。この配列では、それぞれの電球とスイッチがつながっている場合は1、つながっていない場合は0とします。

この配列がつくれればbit探索が容易にできます。
方針は、

1. 上記記載のように電球とスイッチの2次元配列を作成
2. すべてのスイッチについて、onとするかoffとするかをbit全探索（for文のiとjで実現）
3. それぞれのbitの状態で、すべての電球が点灯するか否かをチェック（for文のkで実現）
4. 電球が点灯するか否かのチェックの際はflagで情報を持って置き、flagがFalseとなった段階でbreakする
5. flagがTrueとなる状態数を数えあげてそれが答え
   です。

---

012. AtCoder Beginner Contest 002 D - 派閥

回答

import itertools

N, M = map(int, input().split()) # Nは議員の数、Mは人間関係の数
members = list(range(N))
relation = [[0] * N for _ in range(N)] # 人間関係のマトリクス
for _ in range(M):
    x, y = map(int, input().split())
    x -= 1
    y -= 1
    relation[x][y] = 1
    relation[y][x] = 1

max_answer = 0
for group_num in range(1, N+1): # グループの人数をすべて試す
    temp_answer = 0
    for c in itertools.combinations(members, group_num): # memberがら特定の人数を列挙する
        flag = True
        for base_c in c: # 人間関係をチェックする人
            for check_c in c: # チェック対象となるメンバー
                if base_c == check_c: # 自分自身はチェックしない
                    continue
                if relation[base_c][check_c] == 0: # 人間関係がない場合はすぐにbreak
                    flag = False
                    break
            if not flag:
                break
        if flag: # flagがTrueのときだけOK
            temp_answer = group_num
    max_answer = max(max_answer, temp_answer)

print(max_answer)

itertools を使わないでできなかったのでやむを得なくitertools.combinationsを使用しました。
コードがやや複雑になりましたがやっていることは下記の通り単純です。

1. N人の議員から抽出する任意の人数```group_num``を指定する。これを1~N人まで全通りやります
2. N人のmembersからgroup_num人を抽出する場合分けをすべて列挙します
3. 各議員(base_c)について、上記2で抽出されたメンバー(check_c)と知り合いか否かをチェックします
4. チェックの際は自分自身をチェックしないように気を付け、知り合い出なかった場合はすぐにbreakすることとします
5. すべての議員(base_c)で知り合い関係があると判断された場合(flag==True)はそのgroup_numが答えの候補です
6. 1~6を繰り返しgroup_numの最大値が答え

です。

---

013. JOI 2008 予選 5 - おせんべい

回答

import numpy as np

R, C = map(int, input().split()) # R:縦、C:横
grid = np.array([list(map(int, input().split())) for _ in range(R)])

max_answer = 0
for i in range(2**R):
    copy_grid = grid.copy() # 行列を書き換えるのでコピーを取っておく
    for row in range(R):
        if (i >> row) & 1: # 数の少ない行方向についてbit探索を行う
            copy_grid[row, :] = copy_grid[row, :] ^ 1 # bit演算で0,1を反転

    col_sum = copy_grid.sum(axis=0) # 各列の合計を求めておく
    for col in range(C):
        if  col_sum[col] >= R // 2 + 1: # 1の数が多い箇所はひっくり返す
            copy_grid[:, col] = copy_grid[:, col] ^ 1
    
    answer = R * C - copy_grid.sum()

    max_answer = max(max_answer, answer)

print(max_answer)

ヒントにもある通り、Rの上限が小さいので、Rを全探索を行い、C側で微調整を行います。
方針は、

1. R方向で裏返す箇所を全通り試す
2. R方向で裏返した後にC方向で微調整として、1となっている列が多い箇所だけ裏返す
3. 出来上がった行列の0の数を数える(全体のマス目 - 合計)
4. それぞれの数えた結果の最大をとったものが答え

です。
裏返す操作である0を1、1を0に変換する方法は、bit演算で^1を使います。

---

014. Square869120Contest #4 B - Buildings are Colorful!

回答

import itertools

N, K = map(int, input().split()) # Nは建物の数、K色以上に塗る
A = list(map(int, input().split()))
N_index = list(range(1,N)) # 一番左は固定なので

min_cost = float('inf')
for target_indexes in itertools.combinations(N_index, K-1): # N_indexからK-1個選ぶ方法を全通り試す
    cost = 0
    copy_A = A.copy()
    for i in target_indexes:
        if copy_A[i] <= max(copy_A[:i]):
            diff = max(copy_A[:i]) - copy_A[i] # ターゲットの建物の場合はそれより左側の建物のmaxより1だけ大きくする
            copy_A[i] += diff + 1
            cost += diff + 1

    min_cost = min(min_cost, cost)

print(min_cost)

方針は、

1. 一番左の建物は固定
2. 2番目の建物以降からK-1を選ぶのを全通り試してみる
3. 選んだ建物のindexをtarget_indexes に格納し、各建物の高さを調べる
4. 建物iがそれ以前の建物のmaxの高さ+1になるように、建物iの高さを調整し、調整した分をcostとして集計
5. 上記を全通りやったあとにcostの最小値をとったものが答え

です。

---