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]

Advertisements