# Put Marbles in Bags
**Difficulty:** HARD
[External](https://leetcode.com/problems/put-marbles-in-bags)
Canonical: https://scaleengineer.com/dsa/problems/put-marbles-in-bags
**Patterns:** [Greedy](https://scaleengineer.com/dsa/patterns/greedy)
**Algorithms:** [Sorting](https://scaleengineer.com/algorithms/sorting)
**Data structures:** Array, Heap (Priority Queue)
**Companies:** [Flipkart](https://scaleengineer.com/companies/flipkart), [DE Shaw](https://scaleengineer.com/companies/de-shaw)
---
## 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_.

**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`

# Approaches
## Sorting Partition Costs
This approach leverages a key insight: the total score is determined by the `k-1` locations where the marbles are divided. Each division point, or 'cut', between adjacent marbles `weights[i]` and `weights[i+1]` contributes `weights[i] + weights[i+1]` to the total score. The base score, `weights[0] + weights[n-1]`, is constant for any partition. Therefore, to find the maximum and minimum scores, we only need to find the `k-1` largest and `k-1` smallest partition costs. This can be easily done by calculating all `n-1` possible partition costs and sorting them.
**Time:** O(n log n), where `n` is the number of marbles. Creating the `partitionCosts` array takes `O(n)` time, sorting it takes `O(n log n)`, and summing the required elements takes `O(k)`. The dominant factor is sorting. · **Space:** O(n) to store the `partitionCosts` array. If in-place sorting is used, this is the main space overhead besides the input.
**Pros:** Conceptually simple and easy to implement.; Correctly solves the problem by identifying the core structure.; Efficient enough to pass within the given constraints.
**Cons:** The time complexity is dominated by sorting the entire `partitionCosts` array, which is not strictly necessary.; Uses `O(n)` extra space, which can be improved if `k` is much smaller than `n`.
### Explanation
The algorithm proceeds as follows:
1. First, we recognize that to divide `n` marbles into `k` bags, we must make `k-1` cuts. A cut between `weights[i]` and `weights[i+1]` separates the marbles into two bags, where `weights[i]` is the last element of one bag and `weights[i+1]` is the first element of the next. The cost formula `weights[start] + weights[end]` for a bag means these two weights are added to the total score.
2. The total score can be expressed as `weights[0] + weights[n-1] + sum_of_costs_of_k-1_cuts`.
3. To find the difference between the maximum and minimum scores, the constant term `weights[0] + weights[n-1]` cancels out. The problem reduces to finding `(sum of k-1 largest cut costs) - (sum of k-1 smallest cut costs)`.
4. We create an array, `partitionCosts`, of size `n-1` to store all possible cut costs, where `partitionCosts[i] = weights[i] + weights[i+1]`.
5. We sort this `partitionCosts` array.
6. The sum of the `k-1` smallest costs is the sum of the first `k-1` elements of the sorted array.
7. The sum of the `k-1` largest costs is the sum of the last `k-1` elements of the sorted array.
8. The final result is the difference between these two sums.

```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;
        }

        // There will be k-1 partitions.
        // The cost of a partition at index i is weights[i] + weights[i+1].
        long[] partitionCosts = new long[n - 1];
        for (int i = 0; i < n - 1; i++) {
            partitionCosts[i] = (long)weights[i] + weights[i+1];
        }

        // Sort the partition costs to easily find the smallest and largest k-1 costs.
        Arrays.sort(partitionCosts);

        long minScoreContribution = 0;
        long maxScoreContribution = 0;

        // Sum the k-1 smallest and k-1 largest partition costs.
        for (int i = 0; i < k - 1; i++) {
            minScoreContribution += partitionCosts[i];
            maxScoreContribution += partitionCosts[n - 2 - i];
        }

        return maxScoreContribution - minScoreContribution;
    }
}
```
### Algorithm
- If `k` is 1 or `n`, return 0 as there's only one possible distribution.
- Create an array `partitionCosts` of size `n-1`.
- Iterate from `i = 0` to `n-2`, calculating `partitionCosts[i] = weights[i] + weights[i+1]`.
- Sort the `partitionCosts` array.
- Initialize `minSum` and `maxSum` to 0.
- Iterate from `i = 0` to `k-2`:
    - Add `partitionCosts[i]` to `minSum`.
    - Add `partitionCosts[n-2-i]` to `maxSum`.
- Return `maxSum - minSum`.

## Optimized Selection using Heaps
This approach improves upon the sorting method. Instead of sorting all `n-1` partition costs, we only need to identify the `k-1` smallest and `k-1` largest costs. This is a classic 'Top K' problem, which can be solved efficiently using heaps (Priority Queues). We can find the `k-1` smallest costs using a Max-Heap of size `k-1` and the `k-1` largest costs using a Min-Heap of size `k-1`. This avoids the `O(n log n)` sorting cost, leading to a better time complexity.
**Time:** O(n log k), where `n` is the number of marbles and `k` is the number of bags. We iterate through `n-1` partition costs, and for each, we perform two heap operations. Each heap operation takes `O(log k)` time as the heaps' sizes are capped at `k-1`. · **Space:** O(k) to store the elements in the two heaps. Each heap stores at most `k-1` elements.
**Pros:** More efficient time complexity `O(n log k)` compared to `O(n log n)`.; More efficient space complexity `O(k)` compared to `O(n)`.; Ideal for scenarios where `k` is much smaller than `n`.
**Cons:** Slightly more complex to implement than the sorting approach.; The constant factors involved might make it slightly slower than sorting for small `n` or when `k` is close to `n`.
### Explanation
The core idea remains the same: find the difference between the sum of the `k-1` largest and `k-1` smallest partition costs. The optimization lies in how we find these sums.
1. We iterate through all `n-1` possible partition costs, `weights[i] + weights[i+1]`, in a single pass.
2. To find the `k-1` smallest costs, we use a Max-Heap (a priority queue that keeps the largest element at the top). We add each partition cost to this heap. If the heap's size grows larger than `k-1`, we remove the largest element (`poll()`). After iterating through all `n-1` costs, this heap will contain the `k-1` smallest costs.
3. Similarly, to find the `k-1` largest costs, we use a Min-Heap. We add each cost and if the size exceeds `k-1`, we remove the smallest element. This leaves us with the `k-1` largest costs.
4. Finally, we sum the elements in both heaps and return the difference of the sums. This method is more efficient in both time and space when `k` is significantly smaller than `n`.

```java
import java.util.PriorityQueue;
import java.util.Collections;

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

        // A max-heap to find the k-1 smallest costs.
        PriorityQueue<Long> maxHeapForSmallest = new PriorityQueue<>(Collections.reverseOrder());
        // A min-heap to find the k-1 largest costs.
        PriorityQueue<Long> minHeapForLargest = new PriorityQueue<>();

        for (int i = 0; i < n - 1; i++) {
            long cost = (long)weights[i] + weights[i+1];
            
            // Maintain k-1 smallest costs in the max-heap
            maxHeapForSmallest.offer(cost);
            if (maxHeapForSmallest.size() > k - 1) {
                maxHeapForSmallest.poll();
            }

            // Maintain k-1 largest costs in the min-heap
            minHeapForLargest.offer(cost);
            if (minHeapForLargest.size() > k - 1) {
                minHeapForLargest.poll();
            }
        }

        long minSum = 0;
        while (!maxHeapForSmallest.isEmpty()) {
            minSum += maxHeapForSmallest.poll();
        }

        long maxSum = 0;
        while (!minHeapForLargest.isEmpty()) {
            maxSum += minHeapForLargest.poll();
        }

        return maxSum - minSum;
    }
}
```
### Algorithm
- If `k` is 1 or `n`, return 0.
- Initialize a Max-Heap `maxHeapForSmallest` and a Min-Heap `minHeapForLargest`.
- Iterate from `i = 0` to `n-2`:
    - Calculate `cost = weights[i] + weights[i+1]`.
    - Add `cost` to `maxHeapForSmallest`. If its size exceeds `k-1`, remove the maximum element.
    - Add `cost` to `minHeapForLargest`. If its size exceeds `k-1`, remove the minimum element.
- Sum all elements from `maxHeapForSmallest` to get `minSum`.
- Sum all elements from `minHeapForLargest` to get `maxSum`.
- Return `maxSum - minSum`.

# Solutions
### Java

```java
class Solution {
public
  long putMarbles(int[] weights, int k) {
    int n = weights.length;
    int[] arr = new int[n - 1];
    for (int i = 0; i < n - 1; ++i) {
      arr[i] = weights[i] + weights[i + 1];
    }
    Arrays.sort(arr);
    long ans = 0;
    for (int i = 0; i < k - 1; ++i) {
      ans -= arr[i];
      ans += arr[n - 2 - i];
    }
    return ans;
  }
}

```

### CPP

```cpp
class Solution {
public:
  long long putMarbles(vector<int> &weights, int k) {
    int n = weights.size();
    vector<int> arr(n - 1);
    for (int i = 0; i < n - 1; ++i) {
      arr[i] = weights[i] + weights[i + 1];
    }
    sort(arr.begin(), arr.end());
    long long ans = 0;
    for (int i = 0; i < k - 1; ++i) {
      ans -= arr[i];
      ans += arr[n - 2 - i];
    }
    return ans;
  }
};

```

### Python

```python
class Solution:
    def putMarbles(self, weights: List[int], k: int) -> int: arr = sorted(a + b for a, b in pairwise(weights)) return sum(arr[len(arr) - k + 1:]) - sum(arr[: k - 1])

```