競プロ初中級者100問

JOI2011予選D 1年生

この記事で使うアルゴリズム

動的計画法

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

https://atcoder.jp/contests/joi2011yo/tasks/joi2011yo_d

この問題は動的計画法を用いて解くことができます。

方針

左からn個の数字を使ってmを実現できる組み合わせの数をdp[n][m]として計算していきます。

コード

A[0]を取り出している意図ですが、これはA[0]=0の場合に対処するためです。この操作をせずにdp[0][0]=1としてしまうと、最初のn=0のループの後に、dp[1][0]=2となってしまいます。（一つ目の数字であるにも関わらず、同じ場所dp[0][0]を参照して二度もらってしまう。）

あるいは、DPの更新部分で条件式を

1
if m+A[n] >= 0:
1
if n>0 and m+A[n] >= 0:

のように書き換えても良いでしょう。

1
2
3
4
5
6
7
8
9
10
11
12
13
N = int(input())
A = list(map(int, input().split()))
M = 20
dp = [[0]*(M+1) for _ in range(N-1)]
dp[0][A[0]] = 1
A = A[1:]
for n in range(N-2):
    for m in range(M+1):
        if m-A[n] >= 0:
            dp[n+1][m] += dp[n][m-A[n]]
        if m+A[n] <= M:
            dp[n+1][m] += dp[n][m+A[n]]
print (dp[-1][A[-1]])

記事情報

  • 投稿日：2020年6月8日
  • 最終更新日：2020年6月8日