Leetcode 1664. Ways to Make a Fair Array

Problem statement

https://leetcode.com/problems/ways-to-make-a-fair-array/

Solution

The idea is to calculate prefix sums for odd and even elements. Or we can evaluate cumulative sum of alternating signs, then we can check in O(1) difference of parts when we remove element on index i.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def waysToMakeFair(self, A):
        n = len(A)
        for i in range(1, n, 2): A[i] = -A[i]
        acc = [0] + list(accumulate(A))
        return sum(acc[i+1] + acc[i] == acc[-1] for i in range(n))