2551. Put Marbles in Bags

In the beginning, I was thinking like

[S, X X X X [X X] .... [Y Y].... [Z] .... E]

So only need to care weights[i] and weights[i+1] but also can be weights[i] + weights[i] (single element)

Think in another way, we just put k - 1 blocker into arr and make k subarray.

[S, X X X | X X | X X Y | Y | Y Y Y E]

If 2 blockers' distance == 1 (Y | Y | Y) then it means we select single(Y).

so we just maintain minHeap and maxHeap, and pop the k-1 element then subtract min from max.

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:

        n = len(weights)

        if n == 1:
            return 0

        minheap = []
        maxheap = []

        for i in range(1, len(weights)):

            heappush(minheap, weights[i] + weights[i-1])
            heappush(maxheap, -weights[i] - weights[i-1])


        currentMin = 0
        currentMax = 0
        
        for j in range(k-1):
            currentMin += heappop(minheap)
            currentMax += heappop(maxheap)

        return abs(currentMax) - currentMin