Program to find minimum number of operations to move all balls to each box in Python

Suppose we have a binary string called boxes, where boxes[i] is '0' indicates the ith box is empty, and '1' indicates it contains one ball. In one operation, we can move one ball from a box to an adjacent box. We need to find 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.

Problem Understanding

If the input is boxes = "1101", then the output will be [4, 3, 4, 5] ?

• To put all balls in first box: take from box2 (1 operation) + from box4 (3 operations) = 4 total

• To put all balls in second box: take from box1 (1 operation) + from box4 (2 operations) = 3 total

• To put all balls in third box: take from box1 (2 operations) + from box2 (1 operation) + from box4 (1 operation) = 4 total

• To put all balls in fourth box: take from box1 (3 operations) + from box2 (2 operations) = 5 total

Algorithm Approach

We use a two-pass approach to efficiently calculate operations ?

1. First pass: Calculate initial cost to move all balls to position 0

2. Second pass: Use the relationship between adjacent positions to compute remaining costs

Implementation

def solve(boxes):
    left = 0   # balls to the left of current position
    right = 0  # balls to the right of current position
    
    # Calculate initial distance to move all balls to position 0
    dist = 0
    for i in range(len(boxes)):
        if boxes[i] == "1":
            dist += i
            if i == 0:
                left += 1
            else:
                right += 1
    
    # Build result array
    arr = [dist]
    for i in range(1, len(boxes)):
        # Cost to move to position i = cost to position i-1 + left balls - right balls
        arr.append(arr[i-1] + left - right)
        if boxes[i] == "1":
            left += 1
            right -= 1
    
    return arr

# Test the function
boxes = "1101"
result = solve(boxes)
print(f"Input: {boxes}")
print(f"Output: {result}")

Input: 1101
Output: [4, 3, 4, 5]

How It Works

The algorithm tracks balls on the left and right of each position ?

• Moving right: Left balls get 1 step farther, right balls get 1 step closer

• Formula: cost[i] = cost[i-1] + left_count - right_count

• Time Complexity: O(n) - single pass through the string

• Space Complexity: O(1) - excluding output array

Step-by-Step Example

For boxes = "1101" ?

Conclusion

This solution efficiently calculates minimum operations using the relationship between adjacent positions. The key insight is that moving from position i-1 to i changes the distance for left and right balls predictably.