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

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

今日はEducational DP ContestのB問題を解きます！

B - Frog 2

問題

1,2,…,Nと番号が振られているN個の足場があり、各i(1≤i≤N) について、足場iの高さはh_i。
最初足場1にいるカエルが足場Nまで行きたい。足場iにいるとき、足場i+1,i+2,…,i+Kのどれかへジャンプできて、ジャンプ先の足場をjとすると、コスト|h_i − h_j|を支払う。
カエルが足場Nに辿り着くまでに支払うコストの総和の最小値を求める。

入力(すべて整数)
N：足場の数
K：カエルが何個先まで行けるか
h_i：足場i(1≤i≤N)の高さ

出力(入力がすべて整数なので、これも整数になるはず)
カエルが足場Nに辿り着くまでに支払うコストの総和の最小値

提出→実際の提出

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

dp = [float('inf') for _ in range(n)]
dp[0] = 0


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

print(dp[n-1])

コメントあり

#入力の受け取り
n, k = map(int, input().split())
h = list(map(int, input().split()))

#初期化
dp = [float('inf') for _ in range(n)]
dp[0] = 0

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

#出力
print(dp[n-1])

超絶軽い解説

Aと同じ感じでやったら解けました。
ただ、計算量がo(nk)で、n<10^5, k<100なのでPython3では2104msでTLEを出してしまったため、改めてPyPy3で提出しました(384msだった)。

求めるのは最小値なので、DPのところでテーブルに小さいものをどんどん代入していく必要があります。ということで、0以外のDPテーブルの初期値は無限大(Pythonではfloat('inf')がそれにあたる)にして、始めの比較では必ず代入されるようにしました。

次はC問題を解くぞー！

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