#AtCoder Educational DP Contest / DP まとめコンテスト -B- Python
今回扱うのはこれです
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文を使ったとき
mapとlambda式でごり押した時
PyPy3でfor文を使ったとき
まとめ
- たしかにforは遅いようです。
- よほどの信念でもない限り、ベクトル化は結構ややこしいので素直にPyPy3を使った方がよさそうです。