PythonでJOI難易度5を埋める #27

JOI精進 競プロ

D - 1年生 (A First Grader)

bit全探索だと満点は取れないので，別の方法を考える． i番目までに+か-を入れてきたときの和でDPすればうまくいく．
今の数字を j， i番目に追加する文字を s_iとすると，DP遷移式は，

j - s_i \geq 0のとき，dp[ i + 1][ j - s_i] += dp[ i ][ j ]
j + s_i \leq 20のとき，dp[ i + 1][ j + s_i] += dp[ i ][ j ]

となる．

n = int(input())
s = list(map(int, input().split()))

from pprint import pprint

dp = [[0 for _ in range(21)] for _ in range(n)]
dp[0][s[0]] = 1

for i in range(n - 1):
    for j in range(21):
        tmp1 = j - s[i + 1]
        tmp2 = j + s[i + 1]
        if tmp1 >= 0:
            dp[i + 1][tmp1] += dp[i][j]
        if tmp2 <= 20:
            dp[i + 1][tmp2] += dp[i][j]

print(dp[-2][s[-1]])