[Python] Educational DP Contest A – Frog 2

問題

B – Frog 2

回答

動的計画法を使って解く。

TLEで間に合わなかった回答

in, k = map(int, input().split())
h = [int(i) for i in input().split()]

# dpの最小値を変更する関数
def chmin(a, b):
    if a > b:
        return b
    else:
        return a

# 無限大の値
f_inf = float('inf')

# DP テーブルを初期化　(最小化問題なので INF に初期化)
dp = [f_inf] * (10**5+10)

# 初期条件
dp[0] = 0

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

print(dp[n-1])

他のコードを参考にすると、dpの更新は、forでループを回すのではなく、numpyの行列またはリスト内包表記を使うことで、まとめて計算をすることで、時間を短縮しているようです。

また、以下の記事を参考にして、inputの方法を変更しました。

Pythonの知っておくと良い細かい処理速度の違い8個

numpyを使った回答

import sys
import numpy as np

# inputを高速化する。
input = sys.stdin.readline

n, k = map(int, input().split())
h = [int(i) for i in input().split()]

# dpの配列を作る。
dp = np.zeros(n, dtype=int)
# hの配列を作る。
h = np.array(h)


for i in range(1, n):
    # start は負になることはない。
    start = max(0, i-k)
    # 行列として足し算できるのでforを使う必要がない。
    dp[i] = min(dp[start: i] + np.abs(h[i] - h[start: i]))

print(dp[n-1])