python3で[errror]list index out of rangeが解決できない

Question

atcoder Educational DP Contest B-Frog2をpython3で解いているのですが、

Traceback (most recent call last):
  File ".\atcoder.py", line 20, in <module>
    chmin(dp[i+j], dp[i]+abs(dp[i]-h[i+j]))

IndexError: list index out of range

というエラーがでます。単純なミスだと思うのですが全くわからないまま時間が過ぎてしまったのでどなたか教えていただけると幸いです。コードは以下になります。Qiitaの解説を参考にしています。

import math
import sys
import os
f = open('input.txt', 'r')
sys.stdin = f

def chmin(a, b):
    if a > b:
        a = b
        return a

n, k = map(int, (input().split()))
h = list(input().split())
h = [int(i) for i in h]
dp = [100000]*n
dp[0] = 0
for i in range(0, n):
    for j in range(1, k+1):
        if i+j <= n:
            chmin(dp[i+j], dp[i]+abs(dp[i]-h[i+j]))
print(dp[n-1])

5 3
10 30 40 50 20

Answer 1

ご質問の件は、tachikomaさんが回答済なので別件についてです。

chmin()関数でdp[]の内容を更新したいのだと思いますが、この関数に渡されているa, bはdp[]そのものではなく値のコピーです。なので、a = b と書いてもdp[]の内容を変更することはできません。dp[]を直接更新するようなコードを書く必要があります。

chmin()に渡すbを計算する部分ですが、ここで必要なのは
i番目の足場までのコスト + (i番目の足場からi+j番目の足場に移るためのコスト) なので
dp[i] + abs(dp[i] - h[i+j]) ではなく dp[i] + abs(h[i] - h[i+j]) だと思います。

dpのリストを 100000 で初期化していますが、以下のようにK=1で足場がデコボコで正しい答えが100000を越えるケースでは正しく計算できなくなります。float('inf')とか10**10など、もっと大きな値で初期化する必要があると思います。

100000 1
1 10000 1 10000 1 10000, 1, ...

1-3を反映させたコードは以下の通りです。

def chmin(i, j):
    a, b = dp[i+j], dp[i]+abs(h[i]-h[i+j])
    if a > b:
        dp[i+j] = b


n, k = map(int, (input().split()))
h = list(input().split())
h = [int(i) for i in h]
dp = [10**10]*n
dp[0] = 0
for i in range(0, n):
    for j in range(1, k+1):
        if i+j < n:
            chmin(i, j)
print(dp[n-1])

ただし、上記のコードをPython3 (3.4.3)で提出すると制限時間オーバー(TLE)になります。
chminを別関数にせずにfor jループ内に埋め込むなど、処理時間を速くするための工夫が必要です。
(もしくは、提出言語としてPyPy3を選択するとACになるかもしれません)

Comments

[Sonono]

丁寧にありがとうございます。なんか反映されないなと思って関数をやめたところでした。。。なぜかpythonだとINFないと思い込んでいたので助かりました！ありがとうございます。ご存知でしたら教えていただきたいのですが、質問本文「Qiitaの解説」のリンク先では、c++ですが上のようなchmin()関数で動いています。c++では引数そのものを関数で扱える（変更できる）が、pythonでは値のコピーになってしまうということでしょうか。

[退会済みユーザー]

とりあえずは「pythonでは配列(リスト)の要素を関数に渡しても、呼び出した先の関数で元々の配列を変更することはできない」と覚えておけば良いと思います。←この説明は厳密には正しくないので、気になるようであればpythonの本を読んでみるか「python 参照渡し リスト」あたりのキーワードで検索してみてください。

Answer 2

i+j <= nはi+j < nじゃないとまずそうですね。Pythonのインデックスは0-baseなので、長さnのリスト等に対して有効なインデックスは[0, n-1]なんです。

Comments

[Sonono]

ありがとうございます！エラー解決しました。