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
- if num is non-zero, then
- 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]