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

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]
Updated on: 2021-10-16T10:02:48+05:30

133 Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements