BinarySearch 0790 Ball Moves

Problem statement

https://binarysearch.com/problems/Ball-Moves/

Solution

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

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, A):
        if not A: return []
        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)]