はじめに
こんにちは。これがQiita初投稿のkenです。
この記事を見たということはきっとあなたはDFS全探索を書くのが苦手なんだと思います。私もその気持ちわかりますよ。難しいですよね。
この記事ではPythonを使えばDFS全探索が簡単に書けることもあるよということの紹介をします。
本題
今回はAtcoderからABC233のC - Productを例にとって考えていきたいと思います。
この問題では $\ \prod_{i=1}^N L_i \leq 10^5 \ $という制約がついているため、各袋$i$のなかからひとつずつボールを取り出していく全探索で答えを求めれば良さそうだということがわかります。
ここで一般的(?)に使われるのがDFSを使った全探索で以下のようなコードになります。
import sys
sys.setrecursionlimit(10**7) #再帰上限数を念のため増やしておく
#入力の受け取り
N,X = map(int,input().split())
A = [list(map(int, input().split()))[1:] for _ in range(N)]
ans = 0
def dfs(i,now):
global ans
if i == N:
if now == X:
ans += 1
return
for a in A[i]:
dfs(i+1,now*a)
return
dfs(0,1)
print(ans)
これでACです。
でもDFSってバグらせがちですよね?
再帰ってこれで合ってるのか?って不安になりますよね?
C問題でWA出したくないですよね?
実はPythonならこの問題を再帰を使わず書くことができます。百聞は一見にしかずということで早速コードを見てもらいましょう。
import math
import itertools
N,X = map(int,input().split())
A = [list(map(int, input().split()))[1:] for _ in range(N)]
ans = 0
for v in itertools.product(*A):
if math.prod(v) == X:
ans += 1
print(ans)
どうですか?めっちゃスッキリしてますよね?
鍵となるのはitertools.product()
です。これは引数にリストを入れることでそれらのリストの直積を出力してくれます。引数の*A
というのはアンパックというやつです。リストの中身を取り出して関数に渡してくれています。また、math.prod()
というのはリストなどのイテレータの要素の積を計算するmath
モジュールの関数です。
このようにpythonのライブラリをうまくつかえばDFSを書かなくとも複雑な全探索を書くことができます。
とはいえ
DFSから逃げてるままだといつかきっと壁にぶちあたると思うので用法・用量を守ってお使いください!!最後まで読んでいただきありがとうございました!