#AtCoder Educational DP Contest / DP まとめコンテスト -B- Python

今回扱うのはこれです

atcoder.jp

100点問題なのでそこまで難しくはないので解説はなしです。
ただ、最悪計算量が10^7になります。これはPythonだとわりと厳しい数字です。
あなたがよほどのPython信者でどうしてもPyPyに頼らず、Pythonで通したい場合for文を避けるということで対処します。(Pythonインタープリター方式の言語で、インタープリター方式の言語はfor文が遅いという特徴があるようです。) そういった重い言語を扱う人にとってはfor文を用いるのではなく、ベクトル化して考えることが基本だそうです。
今回の問題ではそういう事情からnumpyを使ってみます。

i番目の足場における理論的に最低のコストをdp[i]に保存します。i番目のコストはdp[i-j]+abs(h[i]-h[i-j]) (ただしi-j>=1)として、jを1からkまでfor文で回すというのが愚直な実装方法ですが、pythonだとこれでTLEします。そこでjに関するfor文をnumpyを使って書き換えます。start = max(1,i-k)として、dp[start:i]+np.abs(h[i]-h[start:i])のminをとればいいことになります。
以下がソースコードです。

import numpy as np

n, k = map(int, input().split())
h = np.array([0] + list(map(int, input().split()))) # index調整
dp = np.zeros(n+1,dtype=np.int)

for i in range(2,n+1):
    start = max(1,i-k)
    dp[i] = min(dp[start:i]+np.abs(h[i]-h[start:i])) # ここがリストなら連結になります
print(dp[n])

おまけですが、mapとlambda式を使って結構強引に書くと下のようになります。でもTLEしたのでC言語API使ってるnumpyの方がやはり早いんでしょうね。

n, k = map(int, input().split())
h = [0] + list(map(int, input().split())) # index調整
dp = [0] * (n+1)

for i in range(2,n+1):
    start = max(1,i-k)
    dp[i] = min(list(map(lambda x,y:x+y, dp[start:i], list(map(lambda x:abs(h[i]-x), h[start:i])))))            
print(dp[n])

実行速度比較
for文を使ったとき

f:id:harutech:20190210202817p:plain
for

mapとlambda式でごり押した時

f:id:harutech:20190210202903p:plain
map,lambda

PyPy3でfor文を使ったとき

f:id:harutech:20190211040221p:plain
PyPy3 for

まとめ

  • たしかにforは遅いようです。
  • よほどの信念でもない限り、ベクトル化は結構ややこしいので素直にPyPy3を使った方がよさそうです。