Put Marbles in Bags — DSA · Engineering Playbook

Problem#

You have weighted marbles in a fixed order and must split them into k non-empty contiguous bags. The cost of a bag is firstWeight + lastWeight. Return the difference between the maximum and minimum total cost over all valid splits.

Examples#

weights = [1,3,5,1], k = 2 → 4
weights = [1, 3], k = 2 → 0

Constraints#

1 <= k <= weights.length <= 10^5, 1 <= weights[i] <= 10^9.

Approach#

Total cost telescopes: weights[0] + weights[n-1] + sum over chosen cuts of (weights[i] + weights[i+1]). The first and last terms are fixed. We choose exactly k - 1 cuts from the n - 1 adjacent pair-sums. Maximum cost takes the top k - 1 sums; minimum takes the bottom k - 1. The answer is their difference.

Solution#

class Solution {
public:
    long long putMarbles(vector<int>& weights, int k) {
        int n = weights.size();
        if (k == 1 || n == k) return 0;
        vector<int> pairs(n - 1);
        for (int i = 0; i < n - 1; ++i) pairs[i] = weights[i] + weights[i + 1];
        sort(pairs.begin(), pairs.end());
        long long lo = 0, hi = 0;
        for (int i = 0; i < k - 1; ++i) {
            lo += pairs[i];
            hi += pairs[n - 2 - i];
        }
        return hi - lo;
    }
};

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

from typing import List

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        n = len(weights)
        if k == 1 or n == k:
            return 0
        pairs = [weights[i] + weights[i + 1] for i in range(n - 1)]
        pairs.sort()
        lo = sum(pairs[: k - 1])
        hi = sum(pairs[-(k - 1):])
        return hi - lo

function putMarbles(weights, k) {
    const n = weights.length;
    if (k === 1 || n === k) return 0n;
    const pairs = new Array(n - 1);
    for (let i = 0; i < n - 1; i++) pairs[i] = weights[i] + weights[i + 1];
    pairs.sort((a, b) => a - b);
    let lo = 0n, hi = 0n;
    for (let i = 0; i < k - 1; i++) {
        lo += BigInt(pairs[i]);
        hi += BigInt(pairs[n - 2 - i]);
    }
    return hi - lo;
}

class Solution {
    public long putMarbles(int[] weights, int k) {
        int n = weights.length;
        if (k == 1 || n == k) return 0L;
        int[] pairs = new int[n - 1];
        for (int i = 0; i < n - 1; i++) pairs[i] = weights[i] + weights[i + 1];
        Arrays.sort(pairs);
        long lo = 0, hi = 0;
        for (int i = 0; i < k - 1; i++) {
            lo += pairs[i];
            hi += pairs[n - 2 - i];
        }
        return hi - lo;
    }
}

function putMarbles(weights: number[], k: number): bigint {
    const n = weights.length;
    if (k === 1 || n === k) return 0n;
    const pairs: number[] = new Array(n - 1);
    for (let i = 0; i < n - 1; i++) pairs[i] = weights[i] + weights[i + 1];
    pairs.sort((a, b) => a - b);
    let lo = 0n, hi = 0n;
    for (let i = 0; i < k - 1; i++) {
        lo += BigInt(pairs[i]);
        hi += BigInt(pairs[n - 2 - i]);
    }
    return hi - lo;
}

Editorial#

The contribution of every internal split between indices i and i+1 is exactly weights[i] + weights[i+1], irrespective of how the rest of the array is partitioned. So the problem decouples into picking the best vs worst k-1 adjacent pair-sums; sort once and pluck both ends.

Complexity#

Time: O(n log n).
Space: O(n).

Concept revision#

Split Array Largest Sum

Problem#

Examples#

Constraints#

Approach#

Solution#

Editorial#

Complexity#

Concept revision#

Related#