Put marbles in bags

Array Greedy Sorting Heap (Priority Queue)

72.2% acceptance

Feb 25, 2026

2647

125

You have k bags. You are given a 0-indexed integer array weights where weights[i] is the weight of the ith marble. You are also given the integer k. Divide the marbles into the k bags according to the following rules: No bag is empty. If the ith marble and jth marble are in a bag, then all marbles with an index between the ith and jth indices should also be in that same bag. If a bag consists of all the marbles with an index from i to j inclusively, then the cost of the bag is weights[i] + weights[j]. The score after distributing the marbles is the sum of the costs of all the k bags. Return the difference between the maximum and minimum scores among marble distributions.

Solution

Rust

Time O(n log n)

Space O(1)

LeetCode

solution.rs

impl Solution {
  pub fn put_marbles(weights: Vec<i32>, k: i32) -> i64 {
    let n = weights.len();
    let k = k as usize;
    if k == 1 || k == n {
      return 0;
    }
    // Each split at position i contributes weights[i] + weights[i+1] to the score
    // (plus weights[0] and weights[n-1] which are always included)
    // We need k-1 splits, so pick k-1 largest for max and k-1 smallest for min
    let mut pair_sums: Vec<i64> = (0..n - 1)
      .map(|i| weights[i] as i64 + weights[i + 1] as i64)
      .collect();
    pair_sums.sort_unstable();
    let m = k - 1;
    let max_score: i64 = pair_sums[n - 1 - m..n - 1].iter().sum();
    let min_score: i64 = pair_sums[..m].iter().sum();
    max_score - min_score
  }
}