Supershipの名畑です。季節の変わり目が訪れるたび、TVアニメ てーきゅうの10期は放送されないのだろうかと期待してしまいます。
はじめに
コードを書きたい気分のときにはよくAtCoderを利用させてもらっています。
AtCoderでは、競技プログラミングの世界で有名なE869120さんが執筆されました書籍「問題解決のための「アルゴリズム×数学」が基礎からしっかり身につく本」の演習問題集が公開されていまして、それをたまに解いています。
そのうちの1問について、解法を残しておきたくなりましたので、残します。
bit全探索と動的計画法のシンプルなサンプルコードとして良いなと思ったのです。
009 - Brute Force 2
問題文は009 - Brute Force 2にございます。
正の整数が書かれたN枚のカードの中から、合計がSになる組み合わせがあるかをYes or Noで答えるという問題です。
1≦N≦20について解けると部分点として500点が、N≦60までを解けると1000点です。
bit全探索による500点のコード
たとえばNが3で、各カードをA0、A1、A2とします。
このときとりうるカードの組み合わせは
- A0
- A1
- A1 + A0
- A2
- A2 + A0
- A2 + A1
- A2 + A1 + A0
の7パターンです。(問題文においてSは1以上のため、カードを選ばないことはない）
2進数で記すなら
- 001
- 010
- 011
- 100
- 101
- 110
- 111
と書けます。左から順にそれぞれの桁がA2、A1、A0の有無を示しているとします。
コードにすると下記です。2進数各桁の取得は右シフトしつつ1との論理積をとっています。1との論理積をとれば下1桁の値になります。
N, S = map(int, input().split())
A = list(map(int, input().split()))
for i in range(1, 2 ** N):
total = 0
for j in range(N):
total += A[j] * ((i >> j) & 1) # 2進数にした各桁が1のときだけAを足す
if total == S:
print('Yes')
exit()
print('No')
計算量がNの2乗に比例するため、Nが大きくなると1000msecを超えてTLE(Time Limit Exceeded)です。
動的計画法(DP)による1000点のコード
たとえばNが2で、各カードをA0、A1とします。
そしてA0が2、A1が3だとします。
1枚目のカードとしてA0を選んだとき、その時点での合計値は2です。選ばなければ0です。
つまり、A0だけしかカードがなかった場合、0と2がとりうる数です。
さらに、A1を選ぶか選ばないかまで考えると
- 0 + 0 = 0
- 0 + 3 = 3
- 2 + 0 = 2
- 2 + 3 = 5
この4パターンです。つまり、A0とA1の2枚のカードしかなかった場合、0、2、3、5の4つがとりうる値となります。
Sがこの中にあればYes、なければNoという考え方です。
これをコードで示すとどうなるでしょう。
N, S = map(int, input().split())
A = list(map(int, input().split()))
dp = [False] * (S + 1) # 0〜Sまでの値をとりうるかを格納する配列
dp[0] = True # 0のみTrueにしておく
for i in range(N): # 0枚目から順番にカードを見ていく
for j in reversed(range(0, S)): # dpの値が途中で変わらないようにreversedしている
if dp[j] and j + A[i] <= S:
dp[j + A[i]] = True
if dp[S]: # Sをとりうるか判定する
print('Yes')
else:
print('No')
計算量はS × Nに比例するため、問題におけるどのテストケースでも50msec以内で終わります。
最後に
考え方次第でこれだけ計算量が変わるのって、本当に面白い。
宣伝
SupershipのQiita Organizationを合わせてご覧いただけますと嬉しいです。他のメンバーの記事も多数あります。
Supershipではプロダクト開発やサービス開発に関わる方を絶賛募集しております。
興味がある方はSupership株式会社 採用サイトよりご確認ください。