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

Key takeaway: An excellent problem to learn problem-solving and step-by-step optimization using loop 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 sum of elements at lower indexes is equal to sum of elements at higher indexes. In other words, equilibrium index of an array is an index i such that sum of elements at indices less than i is equal to 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 sum of elements at lower indexes is equal to 0.
• For i = n - 1, we assume that 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.

Important note: We recommend trying this problem on paper for atleast 15 or 30 minutes. Enjoy problem-solving!

Follow-up questions

• Candidate: For calculating equilibrium index, do we need to include index i in the left or right part? Interviewer: No, 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 at max, there is 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, It may be positive, negative, or zero.
• Candidate: Are the input values unique? Interviewer: No, it 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 two nested loops

Solution idea

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

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

Solution steps

1. We declare variables leftSum and righSum.
2. Now we run a loop from i = 0 to n - 1. Inside loop, we find 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]

4. 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]

5. 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

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

Solution pseudocode

int equilibriumIndex(int A[], int n)
{
    int leftSum
    int 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

def equilibriumIndex(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

We are running nested loops and calculating leftSum and rightSum for every index i. Inside outer loop, summation is a critical operation that will define our time complexity.

• First inner loop run from j = 0 to i - 1. Total summation operation = i
• Second inner loop run from j = i + 1 to n - 1. Total summation operation = n - i - 1
• Total summation operation at each iteration of outer loop = i + n - i - 1 = n - 1
• Here outer loop is running n times, so overall count of summation operation = n(n - 1) = O(n²). Time complexity = O(n²).
• We are using constant extra space. Space complexity = O(1)

Using prefix sum array and single loop

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

Solution idea

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

• If we subtract prefix[i] with value at the current index i, we can easily get the value of leftSum. leftSum = prefix[i] - A[i]
• Now, we should find a way to keep track of sum of elements to the right of current index i. The idea is simple: if we subtract total sum of all array elements with prefix[i], we will get the value of rightSum. We can use a temporary variable totalSum to store sum of all array elements, which is equal to prefix[n]. 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.

Solution pseudocode

int equilibriumIndex(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 equilibriumIndex(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 for calculating prefix sum array and equilibrium index, respectively. Time complexity = Time complexity of calculating prefix sum array + Time complexity of calculating equilibrium Index = O(n) + O(n) = O(n)

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

Efficient in-place solution using single loop

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

Solution idea

Before starting the ith iteration of loop, suppose we know totalSum and value of leftSum. Then at 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 an output. Otherwise, we move to the next iteration.
• Before moving to the next iteration, we update leftSum for next iteration i.e. leftSum = leftSum + A[i].

Solution steps

• We initialize variables totalSum and leftSum with 0.
• We run a loop to store total sum of the array in 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.

Solution pseudocode

int equilibriumIndex(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 equilibriumIndex(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 running two separate loops for calculating total sum and equilibrium index, respectively. Time complexity = Time complexity of calculating total sum + Time complexity of finding equilibrium index = O(n) + O(n) = O(n). Space complexity = O(1), we only use variables to store total, left, and right sum.

Critical ideas to think about!

• Can we solve this problem using hash table?
• Can we solve this problem using two pointers approach?
• Can we try to solve this problem using some other approach?
• If there can be multiple equilibrium indexes, How do we modify above code to return all equilibrium indexes? What would be different corner cases?
• In efficient approach, why are we calculating leftSum 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

• Number of buildings facing the sun
• Remove duplicates from the sorted array
• Chocolate Distribution Problem
• Sort an array of 0s, 1s and 2s
• Find minimum and maximum in an array
• Move zeroes to an end
• Product of Array Except Self
• Sort an array in a waveform

Thanks to Navtosh Kumar for his contribution to the content. Please write in the message below if you find anything incorrect, or if you want to share more insight. Enjoy learning, Enjoy algorithms!