leetcode - 3413. Maximum Coins From K Consecutive Bags

Description

There are an infinite amount of bags on a number line, one bag for each coordinate. Some of these bags contain coins.

You are given a 2D array coins, where coins[i] = [li, ri, ci] denotes that every bag from li to ri contains ci coins.

The segments that coins contain are non-overlapping.

You are also given an integer k.

Return the maximum amount of coins you can obtain by collecting k consecutive bags.

Example 1:

Input: coins = [[8,10,1],[1,3,2],[5,6,4]], k = 4

Output: 10

Explanation:

Selecting bags at positions [3, 4, 5, 6] gives the maximum number of coins: 2 + 0 + 4 + 4 = 10.

Example 2:

Input: coins = [[1,10,3]], k = 2

Output: 6

Explanation:

Selecting bags at positions [1, 2] gives the maximum number of coins: 3 + 3 = 6.

Constraints:

1 <= coins.length <= 10^5
1 <= k <= 10^9
coins[i] == [li, ri, ci]
1 <= li <= ri <= 10^9
1 <= ci <= 1000
The given segments are non-overlapping.

Solution

Spent really long time on this one…

If we don’t have k constrains, then it’s a linear sweep problem. With k bags, we will need to have a prefix sum hash map, and if we have cumulative coins at index i, then with k bags it would be prefix[i] - prefix[i - k].

Because of the constrains, we couldn’t iterate from 1 to 10^9, instead we can calculate prefix[i] - prefix[i - k] at each possible optimal point. The possible optimal points are li and ris, because it would be optimal if we start at li and pick k bags forward, or start at ri and pick k bags backwards.

So when preparing linear sweep array, besides (start, 1) and (end, -1), also add (start + k, 0), (end - k, 0).

Time complexity: $o(coins.len)$
Space complexity: $o(coins.len)$

Code

class Solution:
    def maximumCoins(self, coins: List[List[int]], k: int) -> int:
        change_point = {}
        for start, end, coin in coins:
            change_point[start] = change_point.get(start, 0) + coin
            change_point[end + 1] = change_point.get(end + 1, 0) - coin
            change_point[start + k] = change_point.get(start + k, 0)
            change_point[end + 1 - k] = change_point.get(end + 1 - k, 0)
        prefix_sum = {0: 0}
        prev_i = 0
        cur_sum = 0
        cur_change = 0
        res = 0
        for i in sorted(change_point):
            if i < 1:
                continue
            cur_sum += cur_change * (i - prev_i)
            prefix_sum[i] = cur_sum
            cur_change += change_point[i]
            prev_i = i
            if i - k in prefix_sum:
                res = max(res, prefix_sum[i] - prefix_sum[i - k])
        return res