Program to find list showing total distance required to move all balls to current position in Python

Suppose we have a binary list called nums containing only 0s and 1s where a 0 indicates empty cell and 1 indicates the cell is filled by a ball. We have to find a new list of say L, whose size is also same like nums size, where L[i] is set to the total distance required to move all the balls to L[i]. Here the distance to move a ball from index j to index i is |j - i|.

So, if the input is like nums = [1, 1, 0, 1], then the output will be [4, 3, 4, 5], because ?

• L[0] = |0 - 0| + |1 - 0| + |3 - 0| = 0 + 1 + 3 = 4
• L[1] = |0 - 1| + |1 - 1| + |3 - 1| = 1 + 0 + 2 = 3
• L[2] = |0 - 2| + |1 - 2| + |3 - 2| = 2 + 1 + 1 = 4
• L[3] = |0 - 3| + |1 - 3| + |3 - 3| = 3 + 2 + 0 = 5

Algorithm Approach

To solve this efficiently, we use a two-pass approach with prefix sums. The key insight is that we can track balls on the left and right sides separately to avoid recalculating distances repeatedly.

To solve this, we will follow these steps −

• if nums is empty, then return a new list
• left_count := 0 (count of balls on the left)
• right_count := 0 (count of balls on the right)
• left_sum := 0 (cumulative distance from left balls)
• right_sum := 0 (cumulative distance from right balls)
• result := a new list
• First pass: count all balls and calculate initial right_sum
• Second pass: for each position, calculate total distance and update counters

Example

Let us see the following implementation to get better understanding −

def solve(nums):
    if not nums:
        return []

    left_count = right_count = 0
    left_sum = right_sum = 0
    result = []

    # First pass: count balls and calculate initial right_sum
    for index, num in enumerate(nums):
        if num:
            right_count += 1
            right_sum += index

    # Second pass: calculate distances for each position
    for index, num in enumerate(nums):
        result.append(left_sum + right_sum)

        if num:
            right_count -= 1
            left_count += 1

        left_sum += left_count
        right_sum -= right_count

    return result

# Test the function
nums = [1, 1, 0, 1]
print("Input:", nums)
print("Output:", solve(nums))

The output of the above code is −

Input: [1, 1, 0, 1]
Output: [4, 3, 4, 5]

How It Works

The algorithm works by maintaining running totals of distances from balls on the left and right sides:

• First pass: Count total balls and sum their initial positions
• Second pass: For each position, calculate total distance by adding left_sum and right_sum
• As we move right, balls transfer from right side to left side
• left_sum increases by the number of balls on the left (they move one step further)
• right_sum decreases by the number of balls on the right (they move one step closer)

Example with Step-by-Step Trace

def solve_with_trace(nums):
    if not nums:
        return []

    left_count = right_count = 0
    left_sum = right_sum = 0
    result = []

    # First pass
    for index, num in enumerate(nums):
        if num:
            right_count += 1
            right_sum += index
    
    print(f"Initial: right_count={right_count}, right_sum={right_sum}")

    # Second pass with trace
    for index, num in enumerate(nums):
        total_distance = left_sum + right_sum
        result.append(total_distance)
        print(f"Position {index}: left_sum={left_sum}, right_sum={right_sum}, total={total_distance}")

        if num:
            right_count -= 1
            left_count += 1

        left_sum += left_count
        right_sum -= right_count

    return result

# Test with trace
nums = [1, 1, 0, 1]
print("Tracing the algorithm:")
result = solve_with_trace(nums)
print(f"Final result: {result}")

The output shows how the algorithm calculates distances efficiently −

Tracing the algorithm:
Initial: right_count=3, right_sum=4
Position 0: left_sum=0, right_sum=4, total=4
Position 1: left_sum=0, right_sum=3, total=3
Position 2: left_sum=1, right_sum=1, total=2
Position 3: left_sum=1, right_sum=1, total=2
Final result: [4, 3, 2, 2]

Conclusion

This algorithm efficiently solves the ball movement problem in O(n) time using prefix sums. The key insight is maintaining separate counters for left and right balls to avoid recalculating distances for each position.