Pythonでお手軽にbit全探索を行う方法【itertoolsを使おう!】

AtCoder

「10個のスイッチのある時、このうち5個のスイッチがONになっている時だけ特定の動作をさせたい」

この例のように「2通りの状態を持つ要素が幾つかある」場合を全て調べあげようと思ったら、bit全探索を用いるのがスタンダードです。
(要素が少ない場合はfor文を何重か重ねて書くことも出来ますが、要素の個数が多くなると非現実的なため、bit全探索で書くわけです。)

bit全探索を用いる具体的な問題の例を挙げると、AtCoderの以下の問題が良い例です。

・ABC45C たくさんの数式
・ABC128C Switches
・ABC147C HonestOrUnkind2

さて、この記事ではPythonにおいてbit全探索を使う方法を説明します。
Pythonでbit全探索を書く場合、大まかに

・2bit変換を用いる方法
・itertoolsを使う方法(おすすめ)

上記2つの方法に分けられます。

今回はこのうち、あまり一般的ではないがPythonにおいて楽にbit全探索の状態を調べ上げることが出来る「itertools」の方を説明します。

itertoolsを使う方法

まず以下のコードでitertoolsを呼び出しましょう。

import itertools

実はひとくちにitertoolsといっても、その中に様々な道具があるのですが、今回はitertools.productを使います。

「候補となるオブジェクトの集合」と「何回選ぶか」を決めて、次のように表記しましょう。

itertools.product( “候補となるオブジェクトの集合”, “選ぶ回数” )

これで「繰り返しをさせたいオブジェクト同士の全ての組み合わせパターン」を生成することが出来ます。

例として、0と1を入れた配列を2回リピートした場合にどうなるか見てみましょう。

print(list(itertools.product([0,1], repeat=2)))
#[(0, 0), (0, 1), (1, 0), (1, 1)]

このように、「[0,1]という配列内の2つの要素から1つ選ぶ操作を2回繰り返した」場合の、2*2=4つの組み合わせが生成されていることがわかります。
(※listで変換しないとitertoolsオブジェクトがそのまま記されておかしなことになります)

次はこれを用いてbit全探索をしてみましょう。

冒頭に挙げたスイッチの例を用いて、「10個のスイッチのうち5個がONの時に、答えを+1していく」というコードを書きます。

import itertools
ans=0
for i in itertools.product([0,1], repeat=10):
  if sum(i)==5:
    ans+=1
print(ans)
#252

これはC(10,5)=252(通り)と一致するため、しっかりと機能していることがわかります。

実は2通りの要素に限らず、「3通り,4通り…とある複数の要素の状態を調べ上げる」といった場合にも、候補となるオブジェクトの中身を変えれば全く同じように調べることが出来ます。

itertoolsの方が理解しやすく応用も効くケースが多いので、どっちを学んでも良いよ!という場合はitertoolsがオススメです。

(bit全探索以外には順列を作る問題などの使い所があります)

まだitertoolsを使ったことがないよ!という方は、是非この機会に覚えてみてください!

コメント

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