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|
• L[1] = |0 - 1| + |1 - 1| + |3 - 1|
• L[2] = |0 - 2| + |1 - 2| + |3 - 2|
• L[3] = |0 - 3| + |1 - 3| + |3 - 3|

So, to move all the balls to L[1] we have to move the ball from index 0 to 1 with distance 1 and move ball from index 3 to 1 with distance 2.

To solve this, we will follow these steps −

• if nums is empty, then
  • return a new list
• left_count := 0
• right_count := 0
• left_sum := 0
• right_sum := 0
• result := a new list
• for each index and value num in nums, do
  • if num is non-zero, then
    • right_count := right_count + 1
    • right_sum := right_sum + index
• for each index and value num in nums, do
  • insert (left_sum + right_sum) at the end of result
  • if num is non-zero, then
    • right_count := right_count - 1
    • left_count := left_count + 1
  • left_sum := left_sum + left_count
  • right_sum := right_sum - right_count
• return result

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 = []

   for index, num in enumerate(nums):
      if num:
         right_count += 1
         right_sum += index

   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

nums = [1, 1, 0, 1]
print(solve(nums))

Input

[1, 1, 0, 1]

Output

[4, 3, 4, 5]

Arnab Chakraborty

Updated on: 16-Oct-2021

92 Views

Kickstart Your Career

Get certified by completing the course

Get Started