Solving the Equilibrium Index problem in Python

By John Lekberg on October 03, 2020.

This week's post is about solving the "Equilibrium Index" problem. You will learn:

How to create a brute force solution.
How to create a faster solution by making incremental changes, instead of wastefully recomputing data.

Problem statement

You have an array a of n numbers. E.g.

a = (2, 0, -4, 3, -7, 5, 1)

Your goal is to find the array index, i, such that

sum(left) = sum(right)

Where

left = a₁, ..., a_i-1, and
right = a_i+1, ..., a_n.

E.g.

i = 2

Because

left = (2, 0)
right = (3, -7, 5, 1)
sum(left) = sum(right) = 2

This index i is called the equilibrium index because everything to the left of i is in equilibrium with everything to the right of i.

Sometimes there is no solution. E.g.

a = (5, 3, 2, 0, -7, -4, 1)

has no solution.

How I represent the data in the problem

I represent the array a as a list. E.g.

a = [2, 0, -4, 3, -7, 5, 1]

I represent the equilibrium index i as an int that can slice a. E.g.

i = 2
a[:i]

[2, 0]

a[i + 1 :]

[3, -7, 5, 1]

If no solution exists, I use the value None instead of number.

Creating a brute force solution

A simple way to do a brute force search is to try every index i and check if it is an equilibrium index. Here's a Python function that does that:

def equilibrium_index_bf(a):
    """Solve 'Equilibrium Index' problem with brute force.

    Return equilibrium index i such that

    sum(a[:i]) == sum(a[i+1:])

    Or return None if no solution exists.

    a -- list of numbers.
    """
    for i in range(len(a)):
        sum_left = sum(a[:i])
        sum_right = sum(a[i + 1 :])
        if sum_left == sum_right:
            return i
    return None


equilibrium_index_bf([-7, 1, 5, 2, -4, 3, 0])

equilibrium_index_bf([5, 3, 2, 0, -7, -4, 1])

None

I use the built-in functions range and len to iterate over the indices of a. I use the built-in function sum to sum the slices a[:i] and a[i+1:].

What's the time complexity of this algorithm? For an array a of n numbers,

for i in range(len(a)) executes O(n) times.
- Computing sum_left and sum_right takes O(n) steps.

As a result, the overall time complexity of the algorithm is

O(n²)

Creating a faster solution

How can I improve on the brute force solution?

I don't think that I can do better than

for i in range(len(a)):
    ...

because I need to check every possible equilibrium index.

But I can compute sum_left and sum_right faster by incrementally adding and subtracting. Here's a Python function that does this:

def equilibrium_index_fast(a):
    """Solve 'Equilibrium Index' problem with brute force.

    Return equilibrium index i such that

    sum(a[:i]) == sum(a[i+1:])

    Or return None if no solution exists.

    a -- list of numbers.
    """
    sum_left = 0
    sum_right = sum(a)
    for i in range(len(a)):
        sum_right -= a[i]
        if sum_left == sum_right:
            return i
        sum_left += a[i]
    return None


equilibrium_index_fast([-7, 1, 5, 2, -4, 3, 0])

equilibrium_index_fast([5, 3, 2, 0, -7, -4, 1])

None

What's the time complexity of this algorithm? For an array a of n numbers,

for i in range(len(a)) executes O(n) times.
- Computing sum_left and sum_right takes O(1) steps.

As a result, the overall time complexity of the algorithm is

O(n)

In conclusion...

In this week's post, you learned how to solve the "Equilibrium Index" problem. You learned how to create a brute force solution. And, you learned how to create a faster solution by using incremental changes, instead of recomputing sum_left and sum_right from scratch.

My challenge to you:

Here's a modified problem: instead of splitting the array into two sum-equal parts, try to split the array into three sum-equal parts. This means finding indices i and j such that

sum(X) = sum(Y) = sum(Z)

Where

X = a₁, ..., a_i-1,

Y = a_i+1, ..., a_j-1, and

Z = a_j+1, ..., a_n.

If you enjoyed this week's post, share it with your friends and stay tuned for next week's post. See you then!

(If you spot any errors or typos on this post, contact me via my contact page.)