Minimum Number of Operations to Move All Balls to Each Box

main.md

Question

Approach

I thought of a brute force approach first...

Using this approach, I iterated through each box (i) and computed the total number of operations required to move all the balls to that box. For each box (i), I iterated over all other boxes (j), and then calculated the distance from (j) to ( i), and add the distance to the total if the (j)-th box contains a ball (boxes[j] == '1'). It works even though it's an O(n^2) solution, no TLE today thank God. But this approach won't work for larger constraints, so let's use prefix sum!

Moving to prefix sums

With the help of "sources," let's go over the prefix sum approach. I computed the minimum number of operations for each box using prefix sums and then iterated through the binary string twice (once left-to-right and once right-to-left).

In the left-to-right pass, I tracked the cumulative operations and the count of balls encountered so far. For each box (i), the operations needed to move all the balls from the left were computed and stored.

In the right-to-left pass, I reset the cumulative operations and ball count. This time, I computed the operations required to move all the balls from the right to the current box and added these values to the array calculated in the first pass.

Finally, after completing both passes, the array contained the total number of operations for each box, which I returned as the result.

Implementation

class Solution:
    def minOperations(self, boxes: str) -> List[int]:
        n = len(boxes)
        answer = [0] * n

        operations = 0
        count = 0
        for i in range(n):
            answer[i] = operations
            if boxes[i] == '1':
                count += 1
            operations += count

        operations = 0
        count = 0
        for i in range(n - 1, -1, -1):
            answer[i] += operations
            if boxes[i] == '1':
                count += 1
            operations += count
        
        return answer

Complexities

• Time: O(n) for the left-to-right and right-to-left passes.
• Space: O(n) for the answer array.