Index in an array such that its prefix sum equals its suffix sum - best solution

Question

PROBLEM

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 < Nand the sum of elements of lower indices is equal to the sum of elements of higher indices.

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.

Range of N: [0 ... 100,000].

Elements Range: [−2,147,483,648 ... 2,147,483,647].

Complexity: worst-case time O(N)

MY 5-MIN SOLUTION

This is the intuitive solution by computing the formula performance is O(N^2) as it sums the all array for each iteration and it doesnt work for large entries.

def solution(A):
    result = []
    for i in xrange(len(A)):
        if sum(A[:i]) == sum(A[i+1:]):
            result.append(i)
    if result == []:
        return -1
    return result

BEST SOLUTION

This solution is O(N). Can someone explain the logic behind this solution.

def equi(A):
    result = []
    x=1
    i=1
    r=sum(A)
    for e in A:
        i-=1
        r-=2*e
        if -e==r:
            result.append(-i)
    return result

Answer 1

I believe the solution you posted does not work at all, and it's very difficult to understand anyways. So here's my take on it:

def equi(v):
    left_sum = 0       # sum(v[:p]) at all times.
    right_sum = sum(v) # sum(v[p+1:]) at all times.

    for i in xrange(len(v)):
        right_sum -= v[i]
        if left_sum == right_sum:
            return i   # your p.
        left_sum += v[i]
    return -1 # Because, there may not be an equilibrium index at all.

Basically, instead of recalculating sum(left side) and sum(right side) on every iteration of your loop, you can figure them out with simple math.

State at some point, with an array of size n:

pos1 = i
left_sum1 = v[0] + v[1] + ... + v[i-1]
right_sum1 = v[i+1] + v[i+2] + ... + v[n-1]

Now let's go forward one step and check what we should have:

pos2 = i+1
left_sum2 = v[0] + v[1] + ... + v[i]
right_sum2 = v[i+2] + v[i+2] + ... + v[n-1]

Now, what did change?

left_sum2 - left_sum1 = v[i]
right_sum2 - right_sum1 = -v[i+1]

It may be confusing, but it should be plain to see that there is some way to get the sums of the left and right side at any point by adding and substracting to your previous values.

Answer 2

The O(n) solution is a bit too clever and somewhat obfuscated, but it works just fine. I've rewritten it with meaningful variable names and moved some things around to make it more clear how it works.

def equilibriums(A):               # line 1
    result = []                    # line 2
    difference = sum(A)            # line 3
    
    for p in range(len(A)):        # line 4
        
        difference -= 2*A[p]       # line 5
        
        if difference == -A[p]:    # line 6
            result.append(p)       # line 7
        
    return result                  # line 8

Now an explanation. Let

left  = 0,
right = A[0] + A[1] + ... + A[N-2] + A[N-1] = sum(A), and
difference = right - left = sum(A) - 0 = sum(A)    # <-- explains line 3

When A[0] is removed from right and added to left, difference goes down by 2*A[0]. If A[1] is then moved from right to left, difference goes down by 2*A[1]. Whenever element A[p] is moved, difference goes down by 2*A[p]. That explains lines 4 and 5.

Now, at equilibrium index P, we have:

A[0] + A[1] + ... + A[P−1] = A[P+1] + ... + A[N−2] + A[N−1]                # definition
                           = A[P+1] + ... + A[N−2] + A[N−1] + A[P] - A[P]  # add A[P]-A[P]
                           = A[P] + A[P+1] + ... + A[N−2] + A[N−1] - A[P]  # rearrange
\---- but this = left ---/   \--------- and this = right --------/

or,

left = right - A[P]

and

difference = right - left                 # definition
           = right - (right - A[P])       # substitute
           = A[P]                         # simplify

If we move A[P] from right to left, difference goes down by 2*A[P], and now

difference = A[P] - 2*A[P] = -A[P]

That is, when an equilibrium point is moved from right to left, difference goes from A[P] to -A[P]. So, if difference == -A[P], then P is an equilibrium index. That explains the test in line 6.

Note, left and right aren't really needed for the algorithm. They were just used for the explanation.

equilibriums([1,2,3,0,1,0,0,0,0,1,6])  ==> returns [5, 6, 7, 8]

Answer 3

Another approach would be as follows:

1. Consider what we know at the beginning of our quest - we know the length of the input i.e. len(A), we know the sum of the input i.e. sum(A), we also know that no previous sum has been calculated. Thus these become our initial variables e.g.

   sumA, prevSum, n = sum(A), 0, len(A)
   

2. Next let's consider what happens per iteration. We declare a variable e.g. rem which is the value of (sumA - prevSum -A[i]). In essence at each iteration we want to remove the previous sum(prevSum or leftSum) from the sum of the array and also remove the present value.

Then we check if rem == prevSum. If this is True we return the index, if false we repeat this cycle till all elements in the array have been checked at this point we return a Falsy value.

The code for this is available here