Pythonで最大kステップ以内に最終インデックスへ到達するための最小コストを求めるプログラム

問題の概要

数値のリスト nums と整数 k が与えられているとします。nums[i] の各要素は、インデックス i に着地したときにかかるコストを表しています。プレイヤーはインデックス 0 からスタートし、nums の最後のインデックスを目指します。各ステップでは、現在いる位置 X から、最大 k ステップ先までの任意の位置へジャンプすることができます。

ゴールである最後のインデックスに到達するまでに支払うコストの合計を最小化したい場合、その最小値はいくつになるでしょうか?

入力例

たとえば、nums = [2, 3, 4, 5, 6]、k = 2 という入力が与えられたとします。このとき、出力は 12 になります。これは、インデックス 0・2・4 の位置にある 2 + 4 + 6 を選んで通過することで、合計コスト 12 という最小値が実現できるためです。

解法のアプローチ

この問題は動的計画法(DP)の考え方で解くことができます。各位置 i における最小累積コストは、「直近 k 個以内の位置の中で最小の累積コスト」に nums[i] を加えたものになります。しかし、毎回 k 個の候補をすべて調べると計算量が膨らんでしまうため、優先度付きキュー(ヒープ)を使って「到達可能な範囲内での最小コスト」を効率的に取り出せるようにします。

具体的には、以下の手順で処理を進めます。

• ans := 0 として初期化する
• h := 空のヒープを用意する
• i を 0 から nums のサイズ未満まで繰り返す
  • val := 0 とする
  • h が空でない限り、次を繰り返す
    • [val, index] := ヒープの先頭要素 h[0]
    • index >= i - k であれば(まだジャンプ範囲内なら)、ループを抜ける
    • そうでなければ、範囲外になったのでヒープ h から先頭を削除する
  • ans := nums[i] + val と更新する
  • (ans, i) のペアをヒープ h に挿入する
• 最後に ans を返す

ヒープには「その時点までの累積コスト」と「その位置のインデックス」をペアで格納するため、現在位置から k ステップ以内に届かなくなった古いエントリを自動的に排除しながら、常に有効な最小コストを参照できます。

実装例

それでは、実際のPythonコードを見て理解を深めましょう。

from heapq import heappush, heappop
class Solution:
   def solve(self, nums, k):
      ans = 0
      h = []
      for i in range(len(nums)):
         val = 0
         while h:
            val, index = h[0]
            if index >= i - k:
               break
            else:
               heappop(h)
         ans = nums[i] + val
         heappush(h, (ans, i))
      return ans

ob = Solution()
nums = [2, 3, 4, 5, 6]
k = 2
print(ob.solve(nums, k))

入力

[2, 3, 4, 5, 6], 2

出力

12

まとめ

このプログラムでは、ヒープ(優先度付きキュー)を活用することで、各位置への到達時に「直近 k ステップ以内の最小累積コスト」を高速に取得できます。計算量は O(n log n) となり、すべての経路を総当たりで調べる方法よりも大幅に効率的です。ジャンプ幅の制約があるコスト最小化問題において、ヒープによる範囲管理は非常に有用なテクニックといえます。