Put Marbles in Bags

Problem

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.

Examples

Example 1:

1
2
3
4
5
6

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:

1
2
3
4

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 <= 10^5
• 1 <= weights[i] <= 10^9

Solution

Method 1 - Sorting

First, lets look at dependencies:

• Each bag must contain marbles with consecutive indices.
• Each bag’s “cost” is calculated as the sum of the first and last marble in that bag.

Divide the marbles into k bags such that we can evaluate the maximum and minimum scores achievable and return their difference.

Special Case

There are two distinct cases where only one distribution is possible: k=1 or n=k.

In the case of k=1, all marbles must go into a single bag. For n=k, each marble goes into its own bag. Since there’s only one distribution possible in both scenarios, the difference between the maximum and minimum scores will inevitably be 0.

General Cases

Now let’s consider the general scenario where n marbles need to be divided into k bags. Regardless of the distribution method, the weights of the first and last marbles will always be included in the score. The only variation in the score arises from how the marbles are split in between.

To distribute the marbles into k bags, place k-1 dividers among the n-1 gaps between consecutive marbles. Placing a divider after the i-th marble increases the score by weights[i] + weights[i+1]. For example, with n=10 and k=3, one distribution could be [0,1,2 | 3,4,5,6,7 | 8,9], where placing a divider after marble 2 increases the score by weights[2] + weights[3].

Thus, the problem involves identifying the k-1 smallest or largest values in the n-1 pairwise sums (weights[i] + weights[i+1]) from these gaps. Sorting these sums allows us to pick the k-1 values required to compute the minimum and maximum scores and return their difference.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int:
        n = len(weights)
        
        # Compute pairwise sums
        sums = [weights[i] + weights[i+1] for i in range(n-1)]
        
        # Sort sums
        sums.sort()
        
        # Maximum score
        max_score = sum(sums[-(k-1):]) + weights[0] + weights[-1]
        
        # Minimum score
        min_score = sum(sums[:k-1]) + weights[0] + weights[-1]
        
        # Return the difference
        return max_score - min_score

Complexity

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

Method 2 - Sorting with In-place Array Modification

We can improve upon #Method 1 - Sorting. We can store candidates array in weights array in place.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

public long putMarbles(int[] weights, int k) {
	int n = weights.length;
	if (k == n) {
		return 0;
	}

	for (int i = 1; i < n; i++) {
		weights[i-1] = weights[i - 1] + weights[i];
	}
	weights[n-1] = Integer.MAX_VALUE;
	
	Arrays.sort(weights);

	long ans = 0L;

	for (int i = 0; i < k-1; i++) {
		ans += weights[n-2-i] - weights[i];
	}

	return ans;

}

Complexity

• Time: O(n log n)
• Space: O(1) if we consider array modified as no new array used.