Equilibrium Index of an Array

Difficulty: Easy, Asked-in: Amazon, Adobe, Hike.

Key takeaway: An excellent problem to learn problem-solving and step-by-step optimization using prefix array and single loop using variables.

Let's understand the problem

Write a program to find the equilibrium index of an array. The equilibrium index of an array is an index such that the sum of elements at lower indexes is equal to the sum of elements at higher indexes.

In other words, the equilibrium index of an array is an index 'i' such that the sum of elements at indices less than 'i' is equal to the sum of elements at indices greater than 'i'. A[0] + A[1] + ... + A[i - 1] = A[i + 1] + ... + A[n - 1], where 0 <= i <= n - 1.

Problem note:

For i = 0, we assume that the sum of elements at lower indexes is equal to 0.
For i = n - 1, we assume that the sum of elements at higher indexes is equal to 0.

Example 1:

Input: A[] = [-7, 1, 5, 2, -4, 3, 0], Output: 3

Explanation: 3 is an equilibrium index, i.e., A[0] + A[1] + A[2] = A[4] + A[5] + A[6] = -1.

Example 2:

Input: A[] = [0, 1, 3, -2, -1], Output: 1

Explanation: 1 is an equilibrium index, i.e., A[0] = A[2] + A[3] + A[4] = 0.

Example 3:

Input: A[] = [1, 2, -2, -1], Output: -1

Explanation: There is no such equilibrium index in the array.

Follow-up questions with the interviewer

Candidate: For calculating the equilibrium index, do we need to include index i in the left or right part?Interviewer: No, the element at index i is not included in either part.
Candidate: Is there a possibility of multiple equilibrium indexes? Interviewer: No, it is stated that there is atmost only one equilibrium index in the array. Return the index if it is present; otherwise, return -1.
Candidate: Are the input values only positive? Interviewer: No, they may be positive, negative, or zero.
Candidate: Are the input values unique? Interviewer: No, they can be repeated.

Discussed solution approaches

Brute force approach using nested loops
Using prefix sum array and single loop
Efficient approach using single loop

Brute force approach using nested loops

The basic idea is to explore each index i and find the sum of all elements on the left and right sides, excluding the value at index i. If both left and right sums are equal, the index i would be the equilibrium index.

We can use nested loops to implement this: the outer loop explores each index i, and the inner loop determines whether index i is an equilibrium index or not by calculating left and right sums. We use two extra variables: leftSum and rightSum to store the sum of elements on the left and right sides of the current index.

Solution steps

We declare variables leftSum and rightSum.
Now we run a loop from i = 0 to n - 1. Inside the loop, we find the sum of elements towards the left and right sides of the current index.
We initialize leftSum equal to 0 and run a loop from j = 0 to i - 1 to calculate the leftSum.

for(int j = 0; j < i; j = j + 1)
    leftSum = leftSum + A[j]

We initialize rightSum equal to 0 and run a loop from j = i + 1 to n - 1 to calculate the rightSum.

for(int j = i + 1; j < n; j = j + 1)
    rightSum = rightSum + A[j]

If leftSum and rightSum are equal, the current index i is the equilibrium index and we return i as an output.

if(leftSum == rightSum)
   return i

By the end of outer loop, we return -1 because we did not find any such equilibrium index.

C++ implementation code

int equilibriumIndexArray(int A[], int n) 
{
    int leftSum, rightSum;
    for (int i = 0; i < n; i = i + 1) 
    {
        leftSum = 0;
        for (int j = 0; j < i; j = j + 1)
            leftSum = leftSum + A[j];

        rightSum = 0;
        for (int j = i + 1; j < n; j = j + 1)
            rightSum = rightSum + A[j];

        if (leftSum == rightSum)
            return i;
    }
    return -1;
}

Python implementation code

def equilibriumIndexArray(A, n):
    leftSum = 0
    rightSum = 0
    for i in range(n):
        leftSum = 0
        for j in range(i):
            leftSum = leftSum + A[j]
        
        rightSum = 0
        for j in range(i+1, n):
            rightSum = rightSum + A[j]
     
        if leftSum == rightSum:
            return i 
    return -1

Solution analysis

Inside the outer loop, summation is a critical operation that will define our time complexity.

The first inner loop runs from j = 0 to i-1. The total summation operation is i.
The second inner loop runs from j = i+1 to n-1. The total summation operation is n-i-1.
The total summation operation at each iteration of the outer loop is i + n - i - 1 = n - 1.
As the outer loop runs n times, the overall count of the summation operation is n(n-1) = O(n²). Time complexity = O(n²).
We are using constant extra space. Space complexity = O(1).

Using a prefix sum array and a single loop

Now the critical questions are: Can we improve the time complexity further? Can we do some pre-processing to avoid repeated calculations of leftSum and rightSum at each iteration of the outer loop? Let's think!

Solution idea

Suppose we store the prefix sum of the input array in an extra memory prefix[n], which keeps track of the sum of all elements up to any index i at prefix[i]. Using this prefix array, at each iteration of the outer loop, we can easily calculate leftSum and rightSum in O(1). How? Let's see!

If we subtract the prefix[i] from the value at the current index i, we can easily get the value of leftSum i.e. leftSum = prefix[i] - A[i].
Now, we should find a way to keep track of the sum of elements to the right of the current index i. The idea is simple: if we subtract the total sum of all array elements from the prefix[i], we will get the value of rightSum. We can use a temporary variable totalSum to store the sum of all array elements, which is equal to the prefix[n-1]. rightSum = totalSum - prefix[i].
Now, we compare both leftSum and rightSum to find whether current index is the equilibrium index or not.

Solution steps

We declare prefix sum array of size n.
Now we traverse the array to store the prefix sum.

int prefix[n];
prefix[0] = A[0];

for (int i = 1; i < n; i = i + 1)
    prefix[i] = prefix[i - 1] + A[i];

We initialize variable totalSum to store the total sum of array i.e totalSum = prefix[n - 1].
Now we run loop from i = 0 to n - 1 and calculate values of leftSum and rightSum at each iteration. If leftSum equals rightSum, we return current index i.

for(int i = 0; i < n; i = i + 1)
{
    int leftSum = prefix[i] - A[i]
    int rightSum = totalSum - prefix[i]
    if(leftSum == rightSum)
        return i
}

We return -1 by end of loop because we did not find the equilibrium index.

C++ implementation

int equilibriumIndexArray(int A[], int n) 
{
    int prefix[n];
    prefix[0] = A[0];

    for (int i = 1; i < n; i = i + 1)
        prefix[i] = prefix[i - 1] + A[i];

    int totalSum = prefix[n - 1];

    for (int i = 0; i < n; i = i + 1) 
    {
        int leftSum = prefix[i] - A[i];
        int rightSum = totalSum - prefix[i];

        if (leftSum == rightSum)
            return i;
    }

    return -1;
}

Python implementation

def equilibrium_index_array(A, n):
    prefix = [0] * n
    prefix[0] = A[0]

    for i in range(1, n):
        prefix[i] = prefix[i - 1] + A[i]

    total_sum = prefix[n - 1]

    for i in range(n):
        left_sum = prefix[i] - A[i]
        right_sum = total_sum - prefix[i]

        if left_sum == right_sum:
            return i

    return -1

Solution analysis

We are running two separate loops to calculate the prefix sum array and the equilibrium index, respectively. The time complexity is equal to the time complexity of calculating the prefix sum array plus the time complexity of calculating the equilibrium index, which is O(n) + O(n) = O(n).

The above solution uses extra space to store the prefix sum array. The space complexity is O(n).

Efficient in-place solution using a single loop

The critical question is: Can we improve space complexity further and solve the problem without using a prefix array? Can we track the values of leftSum and rightSum while traversing the array itself? Think!

Solution idea

Before starting the ith iteration of the loop, suppose we know the totalSum and the value of leftSum. Then, at the ith iteration:

We can easily update rightSum, which is equal to totalSum - leftSum - A[i].
Now, we check if leftSum and rightSum are equal or not. If yes, we return index i as output. Otherwise, we move to the next iteration.
Before moving to the next iteration, we update leftSum for the next iteration, i.e., leftSum = leftSum + A[i].

Solution steps

We initialize variables totalSum and leftSum to 0.
We run a loop to store the total sum of the array in the variable totalSum.

for(int i = 0; i < n; i = i + 1)
    totalSum = totalSum + A[i]

Now we iterate through array and keep track of variables leftSum and rightSum. We also check whether current index is the equilibrium index or not.

for(int i = 0; i < n; i = i + 1)
{
    int rightSum = totalSum - leftSum - A[i]
    if(leftSum == rightSum)
        return i
    leftSum = leftSum + A[i]
}

We return -1 by end of loop because we did not find the equilibrium index.

C++ implementation

int equilibriumIndexArray(int A[], int n) 
{
    int totalSum = 0;
    int leftSum = 0;
    
    for (int i = 0; i < n; i = i + 1)
        totalSum = totalSum + A[i];
        
    for (int i = 0; i < n; i = i + 1) 
    {
        int rightSum = totalSum - leftSum - A[i];
        if (leftSum == rightSum)
            return i;
        leftSum = leftSum + A[i];
    }
    
    return -1;
}

Python implementation

def equilibriumIndexArray(A, n):
    totalSum = 0
    leftSum = 0
    for i in range(n):
        totalSum = totalSum + A[i]
        
    for i in range(n):
        rightSum = totalSum - leftSum - A[i]
        if leftSum == rightSum:
            return i
        leftSum = leftSum + A[i]
    return -1

Solution analysis

We are using single loops and variables. So time complexity is O(n) and space complexity is O(1).

Critical idea to think!

Can we solve this problem using a hash table?
Can we solve this problem using the two-pointers approach?
Try to solve this problem using this approach: Suppose we start from the mid index and calculate the left sum and the right sum. If both are equal, we return the index. If the left sum is greater than the right sum, we move the midpoint one position to the left and update the left and right sums. Otherwise, we move the pointer one position to the right and update the left and right sums. Does this approach work when we have negative values in the array?
If there can be multiple equilibrium indexes, how do we modify the above code to return all equilibrium indexes? What would be the different corner cases?
In an efficient approach, why are we calculating the left sum after the comparison if (leftSum == rightSum)?

Comparison of time and space complexities

Brute force approach: Time = O(n²), Space = O(1).
Using prefix sum array: Time = O(n), Space = O(n).
Using single loop: Time = O(n), Space = O(1).

Find Equilibrium Index of an Array

Let's understand the problem

Example 1:

Example 2:

Example 3:

Follow-up questions with the interviewer

Discussed solution approaches

Brute force approach using nested loops

Solution steps

C++ implementation code

Python implementation code

Solution analysis

Using a prefix sum array and a single loop

Solution idea

Solution steps

C++ implementation

Python implementation

Solution analysis

Efficient in-place solution using a single loop

Solution idea

Solution steps

C++ implementation

Python implementation

Solution analysis

Critical idea to think!

Comparison of time and space complexities

Similar coding problems to practice

More from EnjoyAlgorithms

Self-paced Courses and Blogs