蟻本 §2-1-1 部分和問題を Python で

2020年4月30日

蟻本 §2-1-1 部分和問題を Python で

さて,これから蟻本の解説を始めていきたいと思います。
まず初回は蟻本 §2-1-1 部分和問題のコードを Python で書いていきたいと思います。
蟻本は C++ で書かれているので,そのコードをほとんど変えずに Python で書いたものをこれから上げていく予定です。また,AtCoder 版!蟻本 (初級編)を参考にさせてもらい,ここに載っている類題についても上げていく予定です。
それでは,さっそく解いていきます。
まずは問題文から。


§2-1-1 部分和問題


問題文

整数 a1, a2, …, an が与えられます。その中からいくつか選び,その和をちょうど k にできるかを判定しなさい。

制約

  • 1≤n≤20
  • -10^8≤ai≤10^8
  • -10^8≤k≤10^8

入力

入力は以下の形式で標準入力から与えられる。

n
a1 a2 … an
k

出力

可能なら 'Yes’ を,不可能なら 'No’ を出力せよ。


入力例 1

4
1 2 4 7
13

出力例 1

Yes

13 = 2 + 4 + 7 で可能です。

入力例 2

4
1 2 4 7
15

出力例 2

No


解説

深さ優先探索 (DFS) の問題として与えられる最初の問題になります。
これはある状態から遷移を初め,遷移できなくなったら一つ前に戻るという解き方になります。

今回の問題の解答を Python で書くと以下のようになります。
このコードを理解するために,まずは出力例1の全ての遷移を紙にでも書いてください。その最後の遷移(今回は i = n = 4)の時,sum = k = 13 の状態のみ 再起関数の中で True を返します。具体的には,
dfs(0,0) ← dfs(1,0) ← dfs(2,2) ← dfs(3,6) ← dfs(4,13) = True
というように,後ろから順に True が伝播していって最終的に dfs(0,0) で True を返すという流れになります。

自分もそうでしたが,コードだけ与えられてもすぐには理解できませんでした。実際に遷移図を書いて確かめてみると分かりやすくなるかと思います。

解答

n = int(input())
a = list(map(int, input().split()))
k = int(input())

# iまででsumを作って,残りi以降を調べる
def dfs(i, sum):
    # n個決め終わったら,今までの和sumがkと等しいかを返す
    if i == n:
        return sum == k
    
    # a[i]を使わない場合
    if dfs(i+1, sum):
        return True
    
    # a[i]を使う場合
    if dfs(i+1, sum+a[i]):
        return True
    
    # a[i]を使う使わないに拘らずkが作れないのでFalseを返す
    else:
        return False

if dfs(0,0):
    print('Yes')
else:
    print('No')

いかがだったでしょうか。
アルゴリズム自体の考え方については蟻本に載っていますので,そちらを見て頂ければと思います。

この問題の Atcoder での類題も解説していますので,こちらも一緒にご覧ください。
類題はAtCoder 版!蟻本 (初級編)を参考にさせていただいております。

蟻本
蟻本 §2-2-1 部分和問題
蟻本 §2-2-2 Lake Counting

類題
ARC 061 C – たくさんの数式
ABC 079 C – Train Ticket