Minimum Number of Operations to Move All Balls to Each Box

Patrick Leaton
Published: Nov. 27, 2021, 8:09 p.m.
Problem Description

You have n boxes. You are given a binary string boxes of length n, where boxes[i] is '0' if the ith box is empty, and '1' if it contains one ball.

In one operation, you can move one ball from a box to an adjacent box. Box i is adjacent to box j if abs(i - j) == 1. Note that after doing so, there may be more than one ball in some boxes.

Return an array answer of size n, where answer[i] is the minimum number of operations needed to move all the balls to the ith box.

Each answer[i] is calculated considering the initial state of the boxes.

 

Example 1:

Input: boxes = "110"
Output: [1,1,3]
Explanation: The answer for each box is as follows:
1) First box: you will have to move one ball from the second box to the first box in one operation.
2) Second box: you will have to move one ball from the first box to the second box in one operation.
3) Third box: you will have to move one ball from the first box to the third box in two operations, and move one ball from the second box to the third box in one operation.

Example 2:

Input: boxes = "001011"
Output: [11,8,5,4,3,4]

 

Constraints:

  • n == boxes.length
  • 1 <= n <= 2000
  • boxes[i] is either '0' or '1'.

 

The description was taken from https://leetcode.com/problems/minimum-number-of-operations-to-move-all-balls-to-each-box/.

Problem Solution

#O(N) Time, O(N) Space
class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        n = len(boxes)
        prefix = [0] * n
        suffix = [0] * n
        output = []
       
        count = int(boxes[0])
        for i in range(1, n):
            prefix[i] = prefix[i-1] + count
            count += int(boxes[i])
       
        count = int(boxes[-1])
        for i in range(n-2, -1, -1):
            suffix[i] = suffix[i+1] + count
            count += int(boxes[i])
       
        for i in range(n):
            output.append(prefix[i] + suffix[i])
       
        return output

Problem Explanation

Let's first think about how we could solve this problem.

We could iterate through each box, and for each box, we could have two pointers that would traverse each range from the current box's index to the end of the list.

->    <-
0
0
1
0
1
1

If we found a ball, we'd increment the number of operations to move that ball to the current box by the range distance between the ball and the box, as that would be the distance it would have to travel.

This approach would require a traversal through the list, and each iteration would require another traversal through the left and right ranges of the list from the current box.  This would yield a quadratic time complexity.

The problem with this approach is we are recalculating the sums of ranges.

This is where a Prefix Sum would come in handy.

Instead of summing the operations by the distance that each ball would have to travel to get to each current spot, we can instead pass through the array twice, once for each direction.  Through each traversal, we will have a rolling sum of the previous number of operations saved within our prefix/suffix sum array and a previous rolling count of balls that would need to be moved to the current cell.  

When we are finished, we will have two arrays, one for the minimum number of operations to move each ball to the current cell from the left, and one for the minimum number of operations to move each ball to the current cell from the right.

When we sum each index of these two arrays together, we will have the minimum number of operations to move each ball to the current cell from the left and right, we can place these sums together in the output for each index and we have our answer!

Let's start by refactoring the length of the input array to n, since we will be writing this length a lot.  

        n = len(boxes)

 

Next, we will be making our prefix sum array for our left to right traversal, our suffix sum array for our right to left traversal, and our output array where each index will contain the sum of the same indices between our prefix and suffix sum arrays.

        prefix = [0] * n
        suffix = [0] * n
        output = []

 

Since we are going to be looking at the previous value at each index, we will want to start each traversal at the second index from the edge. 

Knowing this, let's initialize our count of previous balls that need to be moved as the value at the first index within the boxes array.

        count = int(boxes[0])

 

Let's start our first traversal.

For each index, we will increment the number of operations within our prefix sum array by the previous number of operations and the previous count of balls to be moved.

Once we have done that, before moving to the next iteration, we will increment the rolling count of balls that need to be moved.

        for i in range(1, n):
            prefix[i] = prefix[i-1] + count
            count += int(boxes[i])

 

We will do the same thing for our suffix sum array moving from right to left.

        count = int(boxes[-1])
        for i in range(n-2, -1, -1):
            suffix[i] = suffix[i+1] + count
            count += int(boxes[i])

 

After we have filled our prefix and suffix sum arrays, we'll need to sum each element between these two arrays to get our answer for each box.

        for i in range(n):
            output.append(prefix[i] + suffix[i])

 

Once we have our answer, we will return it.

        return output