Array     C++     Greedy     Hard     Heap (Priority Queue)     Sorting    

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 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:

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

Solution:

The question amounts to putting up K-1 walls inside the weights array. The score is the sum of all the adjacent members of walls plus weights[0]+weights[n-1]. For example consider the input weights = [a,b,c,d,e,f,g,h,i], K=3 and if we put up walls like this: a b c | d e f g | h i then the score is (a+c) + (d+g) + (h+i) which can also be written as (a+i) + (c+d) + (g+h).

Now we want to find the difference between minimum and maximum scores. Notice that the first and last elements will always be part of any score. So, we need not consider them to calculate the difference. We can use a max-heap and a min-heap to store the sum of all adjacent pairs. In our example, this will be a+b, b+c, c+d, d+e, e+f, f+g, g+h, h+i. Once we have constructed both the heaps, the maximum score is the sum of K-1 largest numbers from max heap and the minimum score is the sum of K-1 smallest numbers from the min heap (plus the first and last, but we are ignoring them because they occur in both). Their difference gives us the answer.

Code

 
class Solution {
public:
    long long putMarbles(vector<int>& weights, int k) 
    {
        int n = weights.size(), cur=weights[0];
        if(k==1 || n==k) return 0;
        priority_queue<int> max_pq;
        priority_queue<int, vector<int>, greater<int>> min_pq;
        for (int i=1; i<n; i++)
        {
            if (i>=2) cur -= weights[i-2];
            cur += weights[i];
            max_pq.push(cur);
            min_pq.push(cur);

        }
        long long res = 0;
        for (int i=0; i<k-1; i++)
        {
            res += max_pq.top();
            res -= min_pq.top();
            max_pq.pop();
            min_pq.pop();
        }
        return res;
    }
};

Complexity

TC: $O(n + k\log(n))$, SC: $O(n)$

If you like this post, please upvote it!