Problem Informations

• Problem Index: 2551
• Problem Link: Put Marbles in Bags
• Topics: Array, Greedy, Sorting, Heap (Priority Queue)

Problem Description

You have k bags. You are given a 0-indexed integer array weights where weights[i] is the weight of the $i^{th}$ 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 $i^{th}$ marble and $j^{th}$ marble are in a bag, then all marbles with an index between the $i^{th}$ and $j^{th}$ 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:

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 \leq k \leq \text{weights.length} \leq 10^5$
• $1 \leq \text{weights}[i] \leq 10^9$

Intuition

The problem requires us to distribute a set of marbles into k bags with specific constraints, aiming to calculate the difference between the maximum and minimum possible scores of such distributions. The constraints ensure that the marbles need to be contiguous within any chosen bag. The cost of each bag depends solely on the weights of the marbles at its boundaries, corresponding to the sum of weights at its starting and ending indices. Thus, meticulously selecting the boundary marbles for the bags is crucial. An intuitive way to tackle this problem is to leverage the pairwise sums of adjacent marbles, which inherently represent potential points to start and end bags (thus defining the cost boundaries).

Approach

The solution involves the following key steps:

1. Pairwise Sum Calculation:

   • Compute the sum of weights for all consecutive marbles. This results in n-1 sums if there are n marbles. Each sum represents a potential partition point between two bags.

2. Sorting:

   • Sort these pairwise sums. The smallest sums represent minimal boundary increases, while the largest sums represent maximal boundary increases. Sorting helps in efficiently selecting the pairwise sums that contribute to either maximizing or minimizing the score.

3. Score Calculation:

   • Maximum Score Calculation: Select the largest (k-1) sums from the sorted list. These sums are added to the score to simulate choosing the largest cost partitions. The choice of (k-1) is based on the fact that the score improves by having (k-1) effective partition points for k bags as the first and last elements of the array contribute once according to the problem definition.
   • Minimum Score Calculation: Similarly, select the smallest (k-1) sums from the sorted list. This results in a minimal score setup, by focusing on small increments in the partition boundaries.

4. Difference Calculation:

   • Compute the final result as the difference between the maximum score and the minimum score calculated in the above steps.

This approach efficiently computes the desired result by exploiting the properties of contiguous partitions, sorted order of pairwise sums, and strategic selection of boundaries. The logic ensures that both maximum and minimum configurations are derived from the same foundational set of potential partitions, resulting in a difference that represents the variability in optimal bagging configurations given the constraints.

Code

C++

class Solution {
public:
    long long putMarbles(vector<int>& weights, int k) {
        int n = weights.size();
        vector<int> splits;

        // Calculate the pair sums of consecutive elements
        for (int i = 0; i < n - 1; i++) {
            splits.emplace_back(weights[i] + weights[i + 1]);
        }

        // Sort the sums of the pairs
        sort(splits.begin(), splits.end());

        // Calculate the maximum score by taking the largest (k-1) sums
        long long maxScore = accumulate(splits.end() - (k - 1), splits.end(), static_cast<long long>(0));

        // Calculate the minimum score by taking the smallest (k-1) sums
        long long minScore = accumulate(splits.begin(), splits.begin() + (k - 1), static_cast<long long>(0));

        // Return the difference between maximum and minimum scores
        return maxScore - minScore;
    }
};

Python

class Solution:
    def putMarbles(self, weights, k):
        n = len(weights)
        splits = []

        # Calculate the pair sums of consecutive elements
        for i in range(n - 1):
            splits.append(weights[i] + weights[i + 1])

        # Sort the sums of the pairs
        splits.sort()

        # Calculate the maximum score by taking the largest (k-1) sums
        maxScore = sum(splits[-(k - 1):])

        # Calculate the minimum score by taking the smallest (k-1) sums
        minScore = sum(splits[:k - 1])

        # Return the difference between maximum and minimum scores
        return maxScore - minScore

Complexity

• Time complexity: $O(n \log n)$
  The time complexity is dominated by the sorting operation on the splits vector, which has $n - 1$ elements. Sorting this vector takes $O((n - 1) \log (n - 1))$ time, which simplifies to $O(n \log n)$ as $n$ becomes large. The subsequent accumulation operations to compute maxScore and minScore each take $O(k)$ time, where $k \leq n$, so these do not affect the overall $O(n \log n)$ time complexity.

• Space complexity: $O(n)$
  The space complexity is determined by the splits vector, which stores $n - 1$ elements, making its space complexity $O(n)$. Additional space is used for variables such as maxScore and minScore, but these require only $O(1)$ space. Thus, the overall space complexity is $O(n)$.