Find Equilibrium Index of an Array

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

Key takeaway: This problem is an excellent opportunity to learn problem-solving and step-by-step optimization using loops and 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

1. We declare variables leftSum and rightSum.
2. 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.
3. 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

3. 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] i.e. 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.

for(int i = 0; i < n; i = i + 1)
{
    if(i == 0)
        prefix[i] = A[i]
    else 
        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];
    for (int i = 0; i < n; i = i + 1) 
    {
        if (i == 0)
            prefix[i] = A[i];
        else
            prefix[i] = A[i] + prefix[i - 1];
    }
    
    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 equilibriumIndexArray(A, n):
    prefix = [0] * n
    for i in range(n):
        if i == 0:
            prefix[i] = A[i]
        else:
            prefix[i] = A[i] + prefix[i-1]
    
    totalSum = prefix[n-1]
    for i in range(n):
        leftSum = prefix[i] - A[i]
        rightSum = totalSum - prefix[i]
        if leftSum == rightSum:
            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

• Can we solve this problem using a hash table?
• Can we solve this problem using the two-pointers approach?
• Can we try to solve this problem using binary search?
• 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)

Suggested coding problems to practice

• Product of Array Except Self
• Find the maximum length of the subarray such that the sum of the elements of the left subarray is equal to the sum of the elements of the right subarray.
• Find maximum equilibrium sum in an array.

If you have any queries/doubts/feedback, please write us at contact@enjoyalgorithms.com. Enjoy learning, Enjoy algorithms!