2551. Put-Marbles-In-Bags

The basic idea is to partition the array into k parts and find the sum of the first and last element of each subarray. We want to do this and find the partition that maximises and minimises this sum. Then we take the difference between the max and min.

So the first step is work out how to partition the array into k parts such that the sum is maximised or minimised.

Initially, my thoughts were to use some kind of dynamic programming:

• consider dp with a parameter of k - no. of bags left
• if k is zero, we just need to grab everything until the end
• if k is not zero, we can either put weights[i] into the current bag (option 1) or start a new bag (option 2)

However, I could not think of a proper recurrence relation for this dp approach, therefore, I consulted the hints.

The hints mentioned that we use priority queues which basically means that I have no idea what I’m doing and I’m failing this technical interview.

Anyway, after looking at the editorials, I realised that the dp recurrence is actually very simple with the subproblem being (x, y) where we split the previous x marbles into y bags and move on to the next larger subproblem of (x + 1, y) or (x, y + 1). However, this has a time complexity of O(n^2) which is not feasible.

The editorial has a nice solution for this, which is to use pairwise sums to calculate min and max score. If we want to partition the array into k parts, it means we need to have k - 1 splits in the array. As the cost for each subarray is calculated only by the two endpoints, the score for each way of partitioning can be calculated just by the sum of the initial element, last element and the sum of the two elements at the split.

Therefore, if we calculate the sum of pairwise elements and sort them, we are able to calculate the min and max score:

• max_score = initial_element + last_element + (the sum of the largest k - 1 pairwise sum)
• min_score = initial_element + last_element + (the sum of the smallest k - 1 pairwise sum)

Therefore, the solution would be:

• max_score - min_score
• (initial_element + last_element + the sum of the largest k - 1 pairwise sum) - (initial_element + last_element + the sum of the smallest k - 1 pairwise sum)
• the sum of the largest k - 1 pairwise sum - the sum of the smallest k - 1 pairwise sum)

class Solution(object):
    def putMarbles(self, weights, k):
        """
        :type weights: List[int]
        :type k: int
        :rtype: int
        """
        n = len(weights)
        pair_weights = [weights[i] + weights[i + 1] for i in range(n - 1)]

        pair_weights.sort()

        answer = 0
        for i in range(k - 1):
            answer += pair_weights[n - 2 - i] - pair_weights[i]

        return answer

Time Complexity: O(n log n)

Space Complexity: O(n)