English
2551. Put Marbles in Bags
Problem Statement
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 andjth
marble are in a bag, then all marbles with an index between theith
andjth
indices should also be in that same bag. - If a bag consists of all the marbles with an index from
i
toj
inclusively, then the cost of the bag isweights[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.
Example 1:
Input: weights = [1,3,5,1], k = 2
Output: 4
Explanation:
The distribution [1],[3,5,1] results in the minimal score of (1+1) + (3+1) = 6.
The distribution [1,3],[5,1], results in the maximal score of (1+3) + (5+1) = 10.
Thus, we return their difference 10 - 6 = 4.
Example 2:
Input: weights = [1, 3], k = 2
Output: 0
Explanation: The only distribution possible is [1],[3].
Since both the maximal and minimal score are the same, we return 0.
Constraints:
1 <= k <= weights.length <= 105
1 <= weights[i] <= 109
Click to open Hints
- Each bag will contain a sub-array.
- Only the endpoints of the sub-array matter.
- Try to use a priority queue.
Solution:
rs
impl Solution {
pub fn put_marbles(weights: Vec<i32>, k: i32) -> i64 {
let n = weights.len();
// pair weights of adjacent bags
let mut pair_weights: Vec<_> = (0..n - 1).map(|i| weights[i] + weights[i + 1]).collect();
// sort the pair weights
pair_weights.sort();
// sum the difference between the largest and smallest pair weights
(0..(k - 1) as usize).fold(0_i64, |ans, i| {
ans + (pair_weights[n - 2 - i] - pair_weights[i]) as i64
})
}
}
...