2551. Put Marbles in Bags


Problem Description

In this problem, you are tasked with distributing a set of marbles, each with a certain weight, into k bags. The distribution must follow certain rules:

  • Contiguity: The marbles in each bag must be a contiguous subsegment of the array. This means if marbles at indices i and j are in the same bag, all marbles between them must also be in that bag.
  • Non-Empty Bags: Each of the k bags must contain at least one marble.
  • Cost of a Bag: The cost of a bag that contains marbles from index i to j (inclusive) is the sum of the weights at the start and end of the bag, that is weights[i] + weights[j].

Your goal is to find the way to distribute the marbles so that you can calculate the maximum and minimum possible scores, where the score is defined as the sum of the costs of all k bags. The result to be returned is the difference between the maximum and minimum scores possible under the distribution rules.

Intuition

To find the maximum and minimum scores, we need to think about what kind of distributions will increase or decrease the cost of a bag.

For the maximum score, you want to maximize the cost of each bag. You can do this by pairing the heaviest marbles together because their sum will be higher. Conversely, for the minimum score, you would pair the lightest marbles together to minimize the sum.

The intuition behind the solution is grounded on the understanding that the cost of a bag can only involve the first and the last marble in it due to the contiguity requirement. Thus, to minimize or maximize the score, you should look at possible pairings of marbles.

A solution approach involves the following steps:

  1. Precompute the potential costs of all possible bags by pairing each marble with the next one, as you will always have to pick at least one boundary marble from each bag.
  2. Sort these potential costs. The smallest potential costs will be at the beginning of the sorted array, and the largest will be towards the end.
  3. To get the minimum score, sum up the smallest k-1 potential costs, which correspond to the minimum possible cost for each of the bags except one.
  4. To get the maximum score, sum up the largest k-1 potential costs, which correspond to the maximum possible cost for each of the bags except one.
  5. The difference between the maximum and minimum scores will provide the result.

The Python solution provided uses the precomputed potential costs and sums the required segments of it to find the maximum and minimum scores, and then computes their difference.

Learn more about Greedy, Sorting and Heap (Priority Queue) patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

Depth first search is equivalent to which of the tree traversal order?

Solution Approach

The implementation of the solution in Python involves the following steps:

  1. Pairwise Costs: Compute the potential costs of making a bag by pairing each marble with its immediate successor. This is done using the expression a + b for a, b in pairwise(weights), which computes the cost for all possible contiguous pairs of marbles and represents the possible costs of bags if they were to contain only two marbles.

  2. Sorting Costs: Once we have all the potential costs, sort them in ascending order using sorted(). This gives us an array where the smallest potential costs are at the start of the array and the largest potential costs are at the end.

  3. Calculating Minimum Score: To obtain the minimum score, sum up the first k-1 costs from the sorted array with sum(arr[: k - 1]). Each of these costs represents the cost of one of the bags in the minimum score distribution, excluding the last bag, because k-1 bags will always have neighboring marbles from the sorted pairwise costs as boundaries, and the last bag will include all remaining marbles which might or might not follow the pairing rule.

  4. Calculating Maximum Score: To get the maximum score, sum up the last k-1 costs from the sorted array with sum(arr[len(arr) - k + 1:]). Each of these costs represents the cost of one of the bags in the maximum score distribution, excluding the last bag, similar to the minimum score calculation but from the other end due to the sorted order.

  5. Computing the Difference: Finally, the difference between the maximum and the minimum scores is computed by return sum(arr[len(arr) - k + 1 :]) - sum(arr[: k - 1]), which subtracts the minimum score from the maximum score to provide the desired output.

This solution uses a greedy approach that ensures the maximum and minimum possible costs by making use of sorted subsegments of potential costs. The data structures used are simple arrays, and the sorting of the pairwise costs array is the central algorithmic pattern that enables the calculation of max and min scores efficiently.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

A heap is a ...?

Example Walkthrough

Let's take a small example to illustrate the solution approach:

Suppose we have an array of marble weights [4, 2, 1, 3] and we want to distribute them into k = 2 bags following the given rules.

Step 1: Pairwise Costs

First, we compute the potential costs for all pairwise marbles:

  • Bag with marbles at index 0 and 1: weights[0] + weights[1] = 4 + 2 = 6
  • Bag with marbles at index 1 and 2: weights[1] + weights[2] = 2 + 1 = 3
  • Bag with marbles at index 2 and 3: weights[2] + weights[3] = 1 + 3 = 4

So the pairwise costs are [6, 3, 4].

Step 2: Sorting Costs

We sort these costs in ascending order: [3, 4, 6].

Step 3: Calculating Minimum Score

To get the minimum score, we need to consider k-1 = 1 smallest pairwise cost, so we sum up the first k-1 costs: 3.

Step 4: Calculating Maximum Score

To get the maximum score, we consider the k-1 largest pairwise costs from the sorted list, so the last k-1 cost: 6.

Step 5: Computing the Difference

The difference between the maximum and minimum scores is 6 - 3 = 3. Therefore, the difference between the highest and the lowest possible scores with the given distribution rules is 3.

In terms of the actual marble distributions:

  • The minimum score distribution would be having one bag with marbles at indices 1 and 2, and the other bag getting the rest (4+2+1+3 = 10, but the last bag's cost is just 4+3=7). Total minimum score: 3 + 7 = 10.
  • The maximum score distribution would be one bag with marbles at indices 0 and 1, and another bag with the rest (2+1+3 = 6, but the last bag's cost is just 2+3=5). Total maximum score: 6 + 5 = 11.

This example clearly illustrates the steps outlined in the solution approach, showing how the precomputed pairwise costs, when sorted and summed appropriately, can yield the minimum and maximum scores, and thus the difference between them.

Solution Implementation

1class Solution:
2    def putMarbles(self, weights, k):
3        n = len(weights)
4        # Initialize dp arrays, dp_max for storing maximum scores, dp_min for storing minimum scores
5        dp_max = [[0] * (k+1) for _ in range(n+1)]
6        dp_min = [[float('inf')] * (k+1) for _ in range(n+1)]
7
8        # Initialize for the case when there is only one bag
9        for i in range(1, n+1):
10            dp_max[i][1] = dp_min[i][1] = weights[0] + weights[i-1]
11
12        # Fill in the dp tables using dynamic programming
13        for i in range(2, n+1):
14            for j in range(2, min(k, i)+1):
15                for m in range(j-1, i):
16                    cost = weights[m] + weights[i-1]  # Calculate the cost for the current grouping
17                    # Update the maximum and minimum scores
18                    dp_max[i][j] = max(dp_max[i][j], dp_max[m][j-1] + cost)
19                    dp_min[i][j] = min(dp_min[i][j], dp_min[m][j-1] + cost)
20
21        # Calculate and return the difference between the maximum and minimum scores
22        return dp_max[n][k] - dp_min[n][k]
23
1public class Solution {
2    public int putMarbles(int[] weights, int k) {
3        int n = weights.length;
4        // Initialize dp arrays for storing maximum and minimum scores
5        int[][] dp_max = new int[n+1][k+1];
6        int[][] dp_min = new int[n+1][k+1];
7
8        // Java doesn't have an equivalent to Python's float('inf'), so we use Integer.MAX_VALUE instead
9        for (int i = 0; i <= n; i++) {
10            for (int j = 0; j <= k; j++) {
11                dp_min[i][j] = Integer.MAX_VALUE;
12            }
13        }
14
15        // Initialize for the case when there is only one bag
16        for (int i = 1; i <= n; i++) {
17            dp_max[i][1] = dp_min[i][1] = weights[0] + weights[i-1];
18        }
19
20        // Fill in the dp tables using dynamic programming
21        for (int i = 2; i <= n; i++) {
22            for (int j = 2; j <= Math.min(k, i); j++) {
23                for (int m = j-1; m < i; m++) {
24                    int cost = weights[m] + weights[i-1];  // Calculate the cost for the current grouping
25                    // Update the maximum and minimum scores
26                    dp_max[i][j] = Math.max(dp_max[i][j], dp_max[m][j-1] + cost);
27                    dp_min[i][j] = Math.min(dp_min[i][j], dp_min[m][j-1] + cost);
28                }
29            }
30        }
31
32        // Calculate and return the difference between the maximum and minimum scores
33        return dp_max[n][k] - dp_min[n][k];
34    }
35}
36
1#include <iostream>
2#include <vector>
3#include <algorithm>
4#include <climits> // For INT_MAX
5
6using namespace std;
7
8class Solution {
9public:
10    int putMarbles(vector<int>& weights, int k) {
11        int n = weights.size();
12        // Initialize dp arrays for storing maximum and minimum scores
13        vector<vector<int>> dp_max(n+1, vector<int>(k+1, 0));
14        vector<vector<int>> dp_min(n+1, vector<int>(k+1, INT_MAX));
15
16        // Initialize for the case when there is only one bag
17        for (int i = 1; i <= n; i++) {
18            dp_max[i][1] = dp_min[i][1] = weights[0] + weights[i-1];
19        }
20
21        // Fill in the dp tables using dynamic programming
22        for (int i = 2; i <= n; i++) {
23            for (int j = 2; j <= min(k, i); j++) {
24                for (int m = j-1; m < i; m++) {
25                    int cost = weights[m] + weights[i-1];  // Calculate the cost for the current grouping
26                    // Update the maximum and minimum scores
27                    dp_max[i][j] = max(dp_max[i][j], dp_max[m][j-1] + cost);
28                    dp_min[i][j] = min(dp_min[i][j], dp_min[m][j-1] + cost);
29                }
30            }
31        }
32
33        // Calculate and return the difference between the maximum and minimum scores
34        return dp_max[n][k] - dp_min[n][k];
35    }
36};
37
1/**
2 * Calculate the difference between the sum of the heaviest and lightest marbles after k turns.
3 *
4 * @param {number[]} weights - An array of integers representing the weights of the marbles.
5 * @param {number} k - The number of turns.
6 * @return {number} The difference in weight between the selected heaviest and lightest marbles.
7 */
8
9function putMarbles(weights: number[], k: number): number {
10    const n: number = weights.length;
11    let dp_max: number[][] = Array.from({ length: n + 1 }, () => Array(k + 1).fill(0));
12    let dp_min: number[][] = Array.from({ length: n + 1 }, () => Array(k + 1).fill(Number.MAX_SAFE_INTEGER));
13
14    // Initialize for the case when there is only one bag
15    for (let i = 1; i <= n; i++) {
16        dp_max[i][1] = dp_min[i][1] = weights[0] + weights[i - 1];
17    }
18
19    // Fill in the dp tables using dynamic programming
20    for (let i = 2; i <= n; i++) {
21        for (let j = 2; j <= Math.min(k, i); j++) {
22            for (let m = j - 1; m < i; m++) {
23                const cost: number = weights[m] + weights[i - 1];  // Calculate the cost for the current grouping
24                // Update the maximum and minimum scores
25                dp_max[i][j] = Math.max(dp_max[i][j], dp_max[m][j - 1] + cost);
26                dp_min[i][j] = Math.min(dp_min[i][j], dp_min[m][j - 1] + cost);
27            }
28        }
29    }
30
31    // Calculate and return the difference between the maximum and minimum scores
32    return dp_max[n][k] - dp_min[n][k];
33}
34
Not Sure What to Study? Take the 2-min Quiz:

Is the following code DFS or BFS?

1void search(Node root) {
2  if (!root) return;
3  visit(root);
4  root.visited = true;
5  for (Node node in root.adjacent) {
6    if (!node.visited) {
7      search(node);
8    }
9  }
10}

Time and Space Complexity

Time Complexity

The time complexity of the provided code can be broken down as follows:

  1. sorted(a + b for a, b in pairwise(weights)):
    • First, the pairwise function creates an iterator over adjacent pairs in the list weights. This is done in O(N) time, where N is the number of elements in weights.
    • The generator expression a + b for a, b in pairwise(weights) computes the sum of each pair, resulting in a total of O(N) operations since it processes each pair once.
    • The sorted function then sorts the sums, which has a time complexity of O(N log N), where N is the number of elements produced by pairwise, which is N - 1. However, we simplify this to O(N log N) for the complexity analysis.

Combining these steps, the time complexity is primarily dominated by the sorting step, resulting in an overall time complexity of O(N log N).

  1. sum(arr[len(arr) - k + 1 :]) and sum(arr[: k - 1]):
    • Both sum(...) operations are performed on slices of the sorted list arr. The number of elements being summed in each case depends on k, but in the worst case, it will sum up to N - 1 elements (the length of arr), which is O(N).

The combination of sorting and summing operations results in a total time complexity of O(N log N), with the sorting step being the most significant factor.

Space Complexity

The space complexity of the code can be examined as follows:

  1. The sorted array arr is a new list created from the sums of pairwise elements, which contains N - 1 elements. Hence, this requires O(N) space.

  2. The slices made in the sum operations do not require additional space beyond the list arr, as slicing in Python does not create a new list but rather a view of the existing list.

Therefore, the overall space complexity of the code is O(N) due to the space required for the sorted list arr.

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

What's the output of running the following function using input [30, 20, 10, 100, 33, 12]?

1def fun(arr: List[int]) -> List[int]:
2    import heapq
3    heapq.heapify(arr)
4    res = []
5    for i in range(3):
6        res.append(heapq.heappop(arr))
7    return res
8
1public static int[] fun(int[] arr) {
2    int[] res = new int[3];
3    PriorityQueue<Integer> heap = new PriorityQueue<>();
4    for (int i = 0; i < arr.length; i++) {
5        heap.add(arr[i]);
6    }
7    for (int i = 0; i < 3; i++) {
8        res[i] = heap.poll();
9    }
10    return res;
11}
12
1class HeapItem {
2    constructor(item, priority = item) {
3        this.item = item;
4        this.priority = priority;
5    }
6}
7
8class MinHeap {
9    constructor() {
10        this.heap = [];
11    }
12
13    push(node) {
14        // insert the new node at the end of the heap array
15        this.heap.push(node);
16        // find the correct position for the new node
17        this.bubble_up();
18    }
19
20    bubble_up() {
21        let index = this.heap.length - 1;
22
23        while (index > 0) {
24            const element = this.heap[index];
25            const parentIndex = Math.floor((index - 1) / 2);
26            const parent = this.heap[parentIndex];
27
28            if (parent.priority <= element.priority) break;
29            // if the parent is bigger than the child then swap the parent and child
30            this.heap[index] = parent;
31            this.heap[parentIndex] = element;
32            index = parentIndex;
33        }
34    }
35
36    pop() {
37        const min = this.heap[0];
38        this.heap[0] = this.heap[this.size() - 1];
39        this.heap.pop();
40        this.bubble_down();
41        return min;
42    }
43
44    bubble_down() {
45        let index = 0;
46        let min = index;
47        const n = this.heap.length;
48
49        while (index < n) {
50            const left = 2 * index + 1;
51            const right = left + 1;
52
53            if (left < n && this.heap[left].priority < this.heap[min].priority) {
54                min = left;
55            }
56            if (right < n && this.heap[right].priority < this.heap[min].priority) {
57                min = right;
58            }
59            if (min === index) break;
60            [this.heap[min], this.heap[index]] = [this.heap[index], this.heap[min]];
61            index = min;
62        }
63    }
64
65    peek() {
66        return this.heap[0];
67    }
68
69    size() {
70        return this.heap.length;
71    }
72}
73
74function fun(arr) {
75    const heap = new MinHeap();
76    for (const x of arr) {
77        heap.push(new HeapItem(x));
78    }
79    const res = [];
80    for (let i = 0; i < 3; i++) {
81        res.push(heap.pop().item);
82    }
83    return res;
84}
85

Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫