Program to find 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