Minimum Cost to Reach Final Index with At Most K Steps in Python

Suppose we have a list of numbers nums and another value k. Here the items at nums[i] represents the costs of landing at index i. If we start from index 0 and end at the last index of nums. In each step we can jump from some position X to any position up to k steps away. We have to minimize the sum of costs to reach last index, so what will be the minimum sum?

So, if the input is like nums = [2, 3, 4, 5, 6] k = 2, then the output will be 12, as we can select 2 + 4 + 6 to get a minimum cost of 12.

To solve this, we will follow these steps −

• ans := 0
• h := an empty heap
• for i in range 0 to size of nums, do
  • val := 0
  • while h is not empty, do
    • [val, index] := h[0]
    • if index >= i - k, then
      • come out from the loop
    • otherwise,
      • delete top from the heap h
  • ans := nums[i] + val
  • insert pair (ans, i) into heap h
• return ans

Let us see the following implementation to get better understanding −

Example 

Live Demo

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))

Input

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

Output

12