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.
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
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
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
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!
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.
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
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
}
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
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.
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!
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!
for(int i = 0; i < n; i = i + 1)
{
if(i == 0)
prefix[i] = A[i]
else
prefix[i] = prefix[i - 1] + A[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
}
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
}
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
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).
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!
Before starting the ith iteration of loop, suppose we know totalSum and value of leftSum. Then at ith iteration:
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]
}
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
}
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
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.
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!
Subscribe to get weekly content on data structure and algorithms, machine learning, system design and oops.