2551. Put Marbles in Bags - Explanation

Problem Link

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)

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

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

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.