経営工学徒の備忘録

~日々の大学生活や留学について~

pythonでビット全探索を実装する【ABC128-C Switches】

スイッチのon,off、何かを含める,含めない、といったように、全部で2のn乗個あるような場合を全て書き出して全探索する方法についてまとめておきます。
例えば以下のような問題です。

atcoder.jp
atcoder.jp
 
本記事では、普段使用しているpythonを用います。

n = int(input())

for i in range(2 ** n):
    for j in range(n):  # ポイント1
        if ((i >> j) & 1):  # ポイント2

ビット全探索の基本的な実装は上記のようになります。
nは、ある・ないを考えるような事象の種類として設定します。

ポイント1
0~2^nまで繰り返されるループの中で、jのループを回します。
つまり、iが0の時jが0,1,2...、iが1の時jが0,1,2....と推移していきます。

ポイント2
最も大事なところです。
ここでは、順に右にシフトさせて最下位bitのチェックを行っています。

もう少し具体的に考えてみましょう。
例えばn=2について調べたいとなった場合、全部で調べる事象は2^2=4通りについてです。
ここで、2進数での表記を考えると、0=00、1=01、2=10、3=11と書くことができます。
この0,1,2,3という値は、最初のループの各iの値に対応していますね。
これらの数字を2進数で表記したものを、j個(ここでは0か1)だけシフトさせて、最下位のbitが0か1であるかというのを((i >> j) & 1)で判定しています。0=off、1=onなどと対応させてあげれば、2種類のスイッチに対してon,offの組み合わせを全部書き出してあげることができました。


最後に、bit全探索を用いてAtcoder Beginner Contest128のC-Switchesを解いてみます。
競技プログラミングは始めたばかりなのでゴリ押しなコード否めないかもしれませんがご了承ください。
atcoder.jp

問題文onとoffの状態を持つ N 個のスイッチと、M個の電球があります。スイッチには1からNの、電球には1からMの番号がついています。電球
iはki個のスイッチに繋がっており、スイッチsi1,si2,...,sikiのうち on になっているスイッチの個数を2で割った余りがpiに等しい時に点灯します。
全ての電球が点灯するようなスイッチの on/off の状態の組み合わせは何通りあるでしょうか。

私はこんなコードを書きました。

N,M = map(int,input().split())
K = [0 for m in range(M)]
count = 0
 
for m in range(M):
  K[m] = list(map(int,input().split()))
P = list(map(int,input().split()))
 
for i in range(2**N):
  op = [0]*N
  for j in range(N):
    if ((i>>j) & 1):
      op[N-1-j] = 1
      
  flag = True
  
  for k in range(M):
    S = 0
    for l in range(K[k][0]):
      S += op[K[k][l+1]-1]
    if S % 2 != P[k]:
      flag = False
  if flag:
    count += 1
    
print(count)