Educational DP Contest B - Frog 2【Pythonの解答例】

前回に続きDPの勉強としてEducational DP Contestに取り組みます。
以前の記事の問題と似ていますので参考に貼っておきます。
https://ebisuke33.hatenablog.com/entry/2020/11/26/225503ebisuke33.hatenablog.com

EDPC B - Frog 2

B - Frog 2

問題文

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

制約

・入力はすべて整数である。
・2≤N≤10^5
・1≤K≤100
・1≤hi≤10^4

解答例

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

dp = [float('inf')] * 100010 # ※1足場数より数が多く、各要素が無限大の配列を準備

dp[0] = 0 # 初期条件

for i in range(n):
    for j in range(1,k+1):
        if i+j <= n-1:
            dpk = dp[i] + abs(h[i] - h[i+j]) # ※2足場iからk離れた足場へ移動したときのコスト
            if dpk < dp[i+j]:
                dp[i+j] = dpk

print(dp[n-1])

解説

DPについてこちらの記事で学習しました。
わかりやすく記載されており大変助かりました。
動的計画法超入門！ Educational DP Contest の A ～ E 問題の解説と類題集 - Qiita
DPの種類として貰うDPと配るDPがあり、以前の記事A - Frogでは貰うDP、今回は配るDPでプログラムが書かれています。
※1では足場の数が制約より多く、各要素が無限大の配列dpを準備しています。要素を無限大としたのは後で出てくるdpkと比較するためです。
B問題ではカエルはk個まで足場を飛ばして移動できるようになりました。
for文は2重ループになっており、※2で足場iからj個離れた足場に移動するときのコストdpkを配っています。
dp[i+j]の要素と配られてきたdpkを比較し、dpkが小さければ配列に代入します。
この繰り返しで、dp[n-1]の最小コストが求まります。

今回のプログラムはPythonとして提出するとTLEが3個出てしまいますが、Pypyで提出するとACとなりました。
まったく知識がなく説明はできないですが、PypyではACとなったことからプログラムの内容は問題なかったと思っています。

いろんなことにつまずきながらですが、今後もEDPCの問題を解いていきます。