2551. Put Marbles in Bags - Explanation

Description

You are given an 0-indexed integer array weights, where weights[i] represents the weight of the i-th marble, and an integer k.

Your task is to divide the marbles into k bags such that:

No bag is empty.
Each bag must contain marbles from a contiguous range of indices. That is, if a bag includes marbles at indices i and j, then all marbles with indices between i and j (inclusive) must also be included in that same bag.
The cost of a bag that includes marbles from index i to j (inclusive) is defined as weights[i] + weights[j].

The total score of a distribution is the sum of the costs of all k bags.

Return the difference between the maximum and minimum possible scores among all valid distributions.

Example 1:

Input: weights = [5,7,9,5], k = 2

Output: 4

Explanation:
The distribution [5,7],[9,5] results in maximal score of (5 + 7) + (9 + 5) = 26.
The distribution [5],[7,9,5] results in minimal score of (5 + 5) + (7 + 5) = 22.
The difference between maximum and minimum scores is 4.

Example 2:

Input: weights = [1,4,2,5,2], k = 3

Output: 3

Explanation:
The distribution [1,4,2],[5],[2] results in maximal score of (1 + 2) + (5 + 5) + (2 + 2) = 17.
The distribution [1],[4],[2,5,2] results in minimal score of (1 + 1) + (4 + 4) + (2 + 2) = 14.
The difference between maximum and minimum scores is 3.

Constraints:

1 <= k <= weights.length <= 100,000
1 <= weights[i] <= 1,000,000,000

Company Tags

Please upgrade to NeetCode Pro to view company tags.

1. Dynamic Programming (Top-Down)

Intuition

We need to partition marbles into k bags and find the difference between maximum and minimum possible scores. Each bag's cost is the sum of its first and last marble weights. Using memoization, we can explore all possible partition points. At each position, we decide whether to make a cut (starting a new bag) or continue the current bag, tracking both max and min scores simultaneously.

Algorithm

Define a recursive function dfs(i, k) that returns [maxScore, minScore] for partitioning marbles from index i with k cuts remaining.
Base cases:
- If k == 0, no more cuts needed, return [0, 0].
- If we reach the end or don't have enough elements for remaining cuts, return invalid values.
At each position, we have two choices:
- Make a partition here: add weights[i] + weights[i + 1] to the score and recurse with k - 1.
- Skip this position: recurse with the same k.
Track both max and min across all choices.
Return maxScore - minScore from the initial call.

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        n = len(weights)
        cache = {}

        def dfs(i, k):
            if (i, k) in cache:
                return cache[(i, k)]
            if k == 0:
                return [0, 0]
            if i == n - 1 or n - i - 1 < k:
                return [-float("inf"), float("inf")]

            res = [0, float("inf")]  # [maxScore, minScore]

            # make partition
            cur = dfs(i + 1, k - 1)
            res[0] = max(res[0], weights[i] + weights[i + 1] + cur[0])
            res[1] = min(res[1], weights[i] + weights[i + 1] + cur[1])

            # skip
            cur = dfs(i + 1, k)
            res[0] = max(res[0], cur[0])
            res[1] = min(res[1], cur[1])

            cache[(i, k)] = res
            return res

        ans = dfs(0, k - 1)
        return ans[0] - ans[1]

Time & Space Complexity

Time complexity: $O(n * k)$
Space complexity: $O(n * k)$

Where $n$ is the number of marbles, and $k$ is the number of bags.

2. Greedy + Sorting

Intuition

The key observation is that the first and last marbles always contribute to the total score regardless of how we partition. What matters is where we place the k-1 dividers. Each divider at position i adds weights[i] + weights[i+1] to the score. To maximize the score, pick the k-1 largest such sums. To minimize, pick the k-1 smallest. The difference between these gives our answer.

Algorithm

If k == 1, return 0 (only one bag, no choices to make).
Create an array of all adjacent pair sums: weights[i] + weights[i+1] for each valid i.
Sort this array.
Sum the smallest k-1 values for the minimum score.
Sum the largest k-1 values for the maximum score.
Return maxScore - minScore.

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        if k == 1:
            return 0

        splits = []
        for i in range(len(weights) - 1):
            splits.append(weights[i] + weights[i + 1])

        splits.sort()
        i = k - 1

        max_score = sum(splits[-i:])
        min_score = sum(splits[:i])
        return max_score - min_score

Time & Space Complexity

Time complexity: $O(n \log n)$
Space complexity: $O(n)$

Where $n$ is the number of marbles, and $k$ is the number of bags.

3. Heap

Intuition

Instead of sorting all adjacent pair sums, we can use two heaps to efficiently track just the k-1 largest and k-1 smallest values. A min-heap keeps the k-1 largest sums (evicting smaller ones), while a max-heap keeps the k-1 smallest sums (evicting larger ones). This is more efficient when k is small relative to n.

Algorithm

If k == 1, return 0.
Initialize a min-heap and a max-heap, both of size at most k-1.
For each adjacent pair sum weights[i] + weights[i+1]:
- Add to the min-heap. If size exceeds k-1, remove the smallest.
- Add to the max-heap. If size exceeds k-1, remove the largest.
Sum all elements in the min-heap for the max score.
Sum all elements in the max-heap for the min score.
Return maxScore - minScore.

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        if k == 1:
            return 0

        max_heap = []
        min_heap = []

        for i in range(len(weights) - 1):
            split = weights[i] + weights[i + 1]

            if len(max_heap) < k - 1:
                heapq.heappush(max_heap, split)
            else:
                heapq.heappushpop(max_heap, split)

            if len(min_heap) < k - 1:
                heapq.heappush(min_heap, -split)
            else:
                heapq.heappushpop(min_heap, -split)

        max_score = sum(max_heap)
        min_score = -sum(min_heap)
        return max_score - min_score

Time & Space Complexity

Time complexity: $O(n \log k)$
Space complexity: $O(k)$

Where $n$ is the number of marbles, and $k$ is the number of bags.