Jump Game VI

  24 Dec 2020

By Wen Xu

heap algorithm dynamic programming

Description

You are given a 0-indexed integer array nums and an integer k.

You are initially standing at index 0. In one move, you can jump at most k steps forward without going outside the boundaries of the array. That is, you can jump from index i to any index in the range [i + 1, min(n - 1, i + k)] inclusive.

You want to reach the last index of the array (index n - 1). Your score is the sum of all nums[j] for each index j you visited in the array.

Return the maximum score you can get.

Example 1:

Input: nums = [1,-1,-2,4,-7,3], k = 2 Output: 7 Explanation: You can choose your jumps forming the subsequence [1,-1,4,3] (underlined above). The sum is 7. Example 2:

Input: nums = [10,-5,-2,4,0,3], k = 3 Output: 17 Explanation: You can choose your jumps forming the subsequence [10,4,3] (underlined above). The sum is 17. Example 3:

Input: nums = [1,-5,-20,4,-1,3,-6,-3], k = 2 Output: 0

Constraints:

1 <= nums.length, k <= 105 -104 <= nums[i] <= 104

Solution

The naive approach would be start from position i, there are at most k choices for the next position (i.e. i + 1, i + 2, …, i + k). So let the maximum score one could obtain from position i be dp[i], then

dp[i] = max(dp[i + 1], dp[i + 2], dp[i + 3], …, dp[i + k]) + nums[i]

one needs to pay attention to the requirement of boundary cases (i + k needs to be smaller than the array size n)

Below is the naive solution


def maxResult(self, nums: List[int], k: int) -> int:
        '''
        let dp[i][k] denotes the maximum score could get starting from position i, then
        dp[i] = max(dp[i + 1], ... dp[i + k]) + nums[i]
        '''
				
        n = len(nums)
        dp = [float('-inf') for i in range(n)]
        dp[n - 1] = nums[n - 1]
        for j in range(n - 2, -1, -1):
            for i in range(1, k + 1):
                if j + i < n:
                    dp[j] = max(dp[j], dp[j + i])
            dp[j] += nums[j]
        return dp[0]

The time complexity of the algorithm is O(nk). It does not pass the OJ. So we need to optimize it a little bit.

We notice for the inner loop we don’t need to check all k elements every time to get the maximum. We can keep all the dp[i] and their indices in a heap. Then we can use O(1) to get the maximum, at the same time, everytime we got an new dp[i], we need to add it to the heap and maintain the heap structure. Below is the implementation

def maxResult(self, nums: List[int], k: int) -> int:
        '''
        let dp[i][k] denotes the maximum score could get starting from position i, then
        dp[i] = max(dp[i + 1], ... dp[i + k]) + nums[i]
        '''
        
				n = len(nums)
        dp = [float('-inf') for i in range(n)]
        dp[n - 1] = nums[n - 1]
        hq = []
        heapq.heappush(hq, (-dp[n - 1], n - 1))
        for j in range(n - 2, -1, -1):
            sc, idx = hq[0]
            while idx > j + k:
                heapq.heappop(hq)
                if len(hq) > 0:
                    sc, idx = hq[0]
            dp[j] = -sc + nums[j]
            heapq.heappush(hq, (-dp[j], j))
        return dp[0]
				

The maximum size the heap can reach is n. So the time complexity is O(nlogn).

comments powered by Disqus