Leetcode 1769. Minimum Number of Operations to Move All Balls to Each Box

Problem statement

https://leetcode.com/problems/minimum-number-of-operations-to-move-all-balls-to-each-box/

Solution

Let us keep for each index how many steps we need to make to move all left to this index balls: A1 and right to this index balls: A2. For this we can use cumulative sums of cumulative sums.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def minOperations(self, boxes):
        A = [int(i) for i in boxes]
        A1 = [0] + list(accumulate(accumulate(A)))[:-1]
        A2 = list(accumulate(accumulate(A[::-1])))[::-1][1:] + [0]
        return [x + y for x, y in zip(A1, A2)]