LeetCode 2551 - Put Marbles in Bags

The problem asks us to distribute weights.length marbles into k contiguous bags, where the cost of a bag is defined as the sum of the first and last marble in that bag.

leetcodehardarraygreedysortingheap-(priority-queue)

LeetCode Problem 2551

Difficulty: 🔴 Hard
Topics: Array, Greedy, Sorting, Heap (Priority Queue)

Solution

Problem Understanding

The problem asks us to distribute weights.length marbles into k contiguous bags, where the cost of a bag is defined as the sum of the first and last marble in that bag. We are to calculate the difference between the maximum and minimum possible total scores across all valid distributions.

The input weights is a list of positive integers representing marble weights, and k is the number of bags. The key constraints are that each bag must be contiguous in the original array and no bag can be empty. The expected output is a single integer, the difference between the maximum and minimum total scores.

Important constraints indicate that weights can have up to 10^5 elements, each weight can be up to 10^9, and k can range from 1 to weights.length. This rules out brute-force enumeration of all distributions due to the combinatorial explosion.

Edge cases to note include when k = 1 (the score is always the sum of the first and last element, so difference is 0) and when k = weights.length (each marble is its own bag, again difference is 0).

Approaches

Brute Force Approach:

A brute-force solution would try every way to partition the array into k contiguous subarrays, compute the sum of the first and last element of each subarray, and keep track of the maximum and minimum total scores. This is correct in principle but computationally infeasible because the number of ways to split an array of length n into k parts is combinatorial: C(n-1, k-1), which is enormous for n = 10^5.

Optimal Approach:

The key observation is that the cost of a bag only depends on its first and last marble. Therefore, for k bags, there will be exactly k-1 cuts between marbles. If we define the "pair sum" between two consecutive marbles weights[i] + weights[i+1], then the total score is weights[0] + weights[n-1] + sum of selected (k-1) pair sums. To maximize the score, select the largest k-1 pair sums; to minimize, select the smallest k-1 pair sums. The difference between maximum and minimum scores is therefore the difference between the sum of the largest k-1 pair sums and the sum of the smallest k-1 pair sums.

This approach reduces the problem to sorting or using a heap on n-1 pair sums, which is efficient for large n.

Approach Time Complexity Space Complexity Notes
Brute Force O(C(n-1, k-1)) O(n) Enumerates all possible partitions, impractical for large n
Optimal O(n log n) O(n) Computes pair sums and sorts to select largest/smallest k-1 sums

Algorithm Walkthrough

  1. Compute an array pair_sums where each element is the sum of consecutive weights: pair_sums[i] = weights[i] + weights[i+1].
  2. Sort pair_sums in ascending order.
  3. To get the minimum score, sum the smallest k-1 pair sums and add weights[0] + weights[-1].
  4. To get the maximum score, sum the largest k-1 pair sums and add weights[0] + weights[-1].
  5. Return the difference between the maximum and minimum scores.

Why it works: The score of a bag depends only on the endpoints. Each cut between marbles determines which endpoints contribute to the total score. By selecting the largest or smallest k-1 cuts, we control the contribution to maximize or minimize the score. This guarantees that the difference between maximum and minimum scores is correctly computed.

Python Solution

from typing import List

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        if k == 1:
            return 0
        
        pair_sums = [weights[i] + weights[i+1] for i in range(len(weights) - 1)]
        pair_sums.sort()
        
        min_score = sum(pair_sums[:k-1])
        max_score = sum(pair_sums[-(k-1):])
        
        return max_score - min_score

Explanation: The code first handles the edge case where k=1. It computes the pair_sums array representing all consecutive marble sums. Sorting allows easy selection of the smallest and largest k-1 sums. Finally, it computes the difference between the maximum and minimum scores as required.

Go Solution

import "sort"

func putMarbles(weights []int, k int) int64 {
    n := len(weights)
    if k == 1 {
        return 0
    }
    
    pairSums := make([]int, n-1)
    for i := 0; i < n-1; i++ {
        pairSums[i] = weights[i] + weights[i+1]
    }
    
    sort.Ints(pairSums)
    
    var minSum, maxSum int64
    for i := 0; i < k-1; i++ {
        minSum += int64(pairSums[i])
        maxSum += int64(pairSums[n-2-i])
    }
    
    return maxSum - minSum
}

Go-specific notes: We explicitly use int64 to avoid integer overflow, especially since weights can be up to 10^9. The slice pairSums is used instead of a list, and sorting uses Go's standard library. Edge case k=1 is handled at the start.

Worked Examples

Example 1: weights = [1,3,5,1], k = 2

  1. Compute pair sums: [1+3=4, 3+5=8, 5+1=6] => [4,8,6]
  2. Sort: [4,6,8]
  3. Minimum score: 4 (smallest 1 sum) + endpoints 1 + 1 = 2 => total 6
  4. Maximum score: 8 (largest 1 sum) + endpoints 2 => total 10
  5. Difference: 10 - 6 = 4

Example 2: weights = [1,3], k = 2

  1. Compute pair sums: [1+3=4]
  2. Only one sum; min and max both 4
  3. Difference = 0

Complexity Analysis

Measure Complexity Explanation
Time O(n log n) Computing pair sums is O(n), sorting O(n log n), selecting sums O(k)
Space O(n) Pair sums array of length n-1

The sorting dominates the complexity. The algorithm efficiently handles the constraints of up to 10^5 elements.

Test Cases

# Provided examples
assert Solution().putMarbles([1,3,5,1], 2) == 4  # Example 1
assert Solution().putMarbles([1,3], 2) == 0     # Example 2

# Edge cases
assert Solution().putMarbles([10], 1) == 0      # Single marble
assert Solution().putMarbles([1,2,3,4,5], 5) == 0  # Each marble is a bag

# Larger case
assert Solution().putMarbles([1,3,5,7,9], 3) == 8  # Check intermediate k
assert Solution().putMarbles([5,1,2,4,6,3], 2) == 8  # Random order
Test Why
[1,3,5,1], k=2 Validates basic functionality for small input
[1,3], k=2 Minimal number of marbles and bags, difference should be 0
[10], k=1 Single marble edge case
[1,2,3,4,5], k=5 Each marble in its own bag, difference should be 0
[1,3,5,7,9], k=3 Intermediate number of bags
[5,1,2,4,6,3], k=2 Random weights for generality

Edge Cases

Single marble, k=1: The difference must be 0 because only one bag exists. The implementation handles this with an early return.

Maximum bags, k=n: Each marble is its own bag, so the cost of each bag is the marble itself twice (first + last), but difference is still 0. Sorting and sum logic still correctly returns 0.

Large weights near 10^9: To prevent overflow in Go, sums are explicitly cast to int64. Python handles arbitrary integers natively. The algorithm still correctly computes pair sums and the difference without numeric errors.