Put Marbles in Bags Interview Question for Google

You are given 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.

Example 1:

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:

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.

Naive Approach

A brute-force approach would involve generating all possible distributions of the marbles into k bags, calculating the score for each distribution, and then finding the difference between the maximum and minimum scores. This would involve exploring all combinations of splitting the weights array into k subarrays, which quickly becomes computationally infeasible as the size of weights increases.

Big O Complexity

Time Complexity: O(2^n), where n is the length of the weights array. This is because, in the worst case, we need to explore all possible subsets.
Space Complexity: O(n) in the worst case due to recursive calls.

Optimal Approach

The key insight is that the weights at the extreme ends of the array (i.e., weights[0] and weights[n-1]) will always be included in the total score, regardless of how the marbles are distributed into bags. The goal is to minimize and maximize the contribution of the weights at the split points.

To minimize the score, we want to minimize the sum of the weights at the k-1 split points. Similarly, to maximize the score, we want to maximize the sum of the weights at the k-1 split points.

Calculate the sum of weights[i] + weights[i+1] for all adjacent pairs of marbles. Store these sums in an array called pairSums.
Sort the pairSums array.
To minimize the score, sum the smallest k-1 values in pairSums. Add weights[0] + weights[n-1] to get the minimum score.
To maximize the score, sum the largest k-1 values in pairSums. Add weights[0] + weights[n-1] to get the maximum score.
Return the difference between the maximum and minimum scores.

Edge Cases

If k is equal to the length of weights, each marble is in its own bag. The minimum and maximum scores will be the same, and the difference is 0.

Big O Complexity

Time Complexity: O(n log n), where n is the length of the weights array. This is due to the sorting step.
Space Complexity: O(n) to store the pairSums array.

Code (Java)

import java.util.Arrays;

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

        long[] pairSums = new long[n - 1];
        for (int i = 0; i < n - 1; i++) {
            pairSums[i] = (long)weights[i] + weights[i + 1];
        }

        Arrays.sort(pairSums);

        long minScore = 0;
        long maxScore = 0;

        for (int i = 0; i < k - 1; i++) {
            minScore += pairSums[i];
            maxScore += pairSums[n - 2 - i];
        }

        return maxScore - minScore;
    }
}