Equilibrium Index

A zero-indexed array A consisting of N integers is given. An equilibrium index of this array is any integer P such that 0 ≤ P < N and the sum of elements of lower indices is equal to the sum of elements of higher indices, i.e.

A[0] + A[1] + ... + A[P−1] = A[P+1] + ... + A[N−2] + A[N−1].

Sum of zero elements is assumed to be equal to 0. This can happen if P = 0 or if P = N−1.

For example, consider the following array A consisting of N = 8 elements:

A[0] = -1 
A[1] = 3 
A[2] = -4
A[3] = 5 
A[4] = 1 
A[5] = -6 
A[6] = 2 
A[7] = 1

P = 1 is an equilibrium index of this array, because:
A[0] = −1 = A[2] + A[3] + A[4] + A[5] + A[6] + A[7]

P = 3 is an equilibrium index of this array, because:
A[0] + A[1] + A[2] = −2 = A[4] + A[5] + A[6] + A[7]

P = 7 is also an equilibrium index, because:
A[0] + A[1] + A[2] + A[3] + A[4] + A[5] + A[6] = 0

P = 8 is not an equilibrium index, because it does not fulfill the condition 0 ≤ P < N.

Write a function that, given a zero-indexed array A consisting of N integers, returns any of its equilibrium indices. The function should return −1 if no equilibrium index exists.

For example, given array A shown above, the function may return 1, 3 or 7, as explained above.

The Idea: Create an integral sum. Then we can check any partial sum in O(1) time by doing sum(ar, a, b) = F[b] - F[a-1]

Complexity: O(N) time and O(N) space

def equindex(ar):
    equi_ar = []
    if not ar:
        return equi_ar

    # create cumulative array
    cum_ar = [ar[0]]
    for i in range(1, len(ar)):
        cum_ar.append(cum_ar[-1]+ar[i])

    # left edge
    sum_left = 0
    sum_right = cum_ar[-1] - cum_ar[0]
    if sum_left == sum_right:
        equi_ar.append(0)

    # inbetween
    for i in range(1, len(ar) - 1):
        sum_left = cum_ar[i-1]
        sum_right = cum_ar[-1] - cum_ar[i]
        if sum_left == sum_right:
            equi_ar.append(i)

    # right edge
    sum_left = cum_ar[-2]
    sum_right = 0
    if sum_left == sum_right:
        equi_ar.append(len(ar) - 1)

    return equi_ar



d = ([-1, 3, -4, 5, 1, -6, 2, 1],
     [-7, 1, 5, 2, -4, 3, 0],
     [2, 4, 6],
     [2, 9, 2],
     [1, -1, 1, -1, 1, -1, 1])

for test in d:
    print(equindex(test))