DPができるようになりたい Part3

こんばんは、ぱんと申します。

今日はEducational DP ContestのC問題を解きます。

問題

C - Vacation

太郎君は、N日の夏休みに、毎日1つ次の活動のうちひとつを選んで行う。

A: 海で泳ぐ。幸福度a_iを得る。
B: 山で虫取りをする。幸福度b_iを得る。
C: 家で宿題をする。幸福度c_iを得る。

太郎君は飽き性なので、2日以上連続で同じ活動を行うことはできない。太郎君が得る幸福度の総和の最大値を求める。

入力(すべて整数)
N：夏休みの日数
a_i, b_i, c_i(1≤i≤N)：各行動を行ったときに得られる幸福度

出力(整数になるはず)
太郎君が得る幸福度の総和の最大値

提出→実際の提出

n = int(input())
h = [] #happiness
for i in range(n):
 a_input, b_input, c_input = map(int, input().split())
 h.append([a_input, b_input, c_input])

dp = [[0 for _ in range(3)] for _ in range(n+1)]


for i in range(1, n+1):
 for j in range(3):
   for k in range(3):
     if j != k and dp[i][j] < dp[i-1][k] + h[i-1][j]:
       dp[i][j] = dp[i-1][k] + h[i-1][j]

ans = 0
for j in range(3):
 if ans < dp[n][j]:
   ans = dp[n][j]

print(ans)

コメントあり

#入力
n = int(input())
h = [] #happiness
for i in range(n):
 a_input, b_input, c_input = map(int, input().split())
 h.append([a_input, b_input, c_input]) #hは二次元配列になる

#初期化
dp = [[0 for _ in range(3)] for _ in range(n+1)]

#DP
for i in range(1, n+1):
 for j in range(3):
   for k in range(3):
     if j != k and dp[i][j] < dp[i-1][k] + h[i-1][j]:
       dp[i][j] = dp[i-1][k] + h[i-1][j]

#出力
ans = 0
for j in range(3):
 if ans < dp[n][j]:
   ans = dp[n][j]
print(ans)

解説

まず幸福値を二次元配列h(happinessの頭文字なので)として受け取ります。つまり、

h = [[a_1, b_1, c_1], [a_2, b_2, c_2], ……, [a_n, b_n, c_n]]

のような感じにします。これがPythonだと

h = []
h.append([a_input, b_input, c_input])

のようにすれば二次元配列になるんですね～～便利。

DPテーブルは

dp[i][j]：i日目終了時点での幸福値の最大値
※ただしi日目は選択肢jを選んだ
※選択肢j：j=0ならa, j=1ならb, j=2ならcの意

というふうにしました。つまり、0日目終了時の幸福値である

dp[0][j] = 0(jに関わらず)

であり、求めたいのは

max(dp[n][j]) (1≤j≤3)

ということになります。

DPテーブルの更新の仕方としては、

i日目終了時の幸福値はdp[i][j] (1≤j≤3)なので、
i-1日目にはjと一致しない選択肢k (1≤k≤3)を選んだはず

と考えて下のようにします。具体的な更新は太線部です(分かりづらいかな…)

#DP
for i in range(1, n+1):
 for j in range(3):
   for k in range(3):
     if j != k and dp[i][j] < dp[i-1][k] + h[i-1][j]:
       dp[i][j] = dp[i-1][k] + h[i-1][j]

なにはともあれ、今回もPyPy3でACできたので良かったです。でも今回で676msだったので、いずれJavaを使う時が来そう。

最後まで読んでいただきありがとうございます。またどうぞ。