Friday, June 6, 2014

Finding the Equilibrium index of an array

I wanted to do a brain teaser today so i took up an algorithmic question to give a shot at. Now i know this question already has many answer on the internet if you do a quick search. But i wanted to try it out with the solution that came to my mind. The problem statement is as follows;

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. For example, in an array A:

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

3 is an equilibrium index, because:
A[0] + A[1] + A[2] = A[4] + A[5] + A[6]

6 is also an equilibrium index, because sum of zero elements is zero, i.e., A[0] + A[1] + A[2] + A[3] + A[4] + A[5]=0

7 is not an equilibrium index, because it is not a valid index of array A.

Write a function int equilibrium(int[] arr); that given a sequence arr[] of size n, returns an equilibrium index (if any) or -1 if no equilibrium indexes exist.

The solution i came up with is as follows;


 
public static int solution(int[] A) {

  if (A == null || A.length < 3)
   throw new RuntimeException("Cannot find equilirbium");
  int pointer = 1;
  int lowerIndCount = A[0];
  int upperIndCount = 0;

  for (int i = 2; i < A.length; i++) {
   upperIndCount += A[i];
   if (lowerIndCount < 0) {
    if (upperIndCount > lowerIndCount && i != A.length - 1)
     continue;
    if (upperIndCount == lowerIndCount && i == A.length - 1)
     break;
    lowerIndCount += A[pointer];
    upperIndCount = 0;
    ++pointer;
    i = pointer;
   } else {
    if (upperIndCount > lowerIndCount) {
     lowerIndCount += A[pointer];
     upperIndCount = 0;
     ++pointer;
     i = pointer;
    }
   }
  }

  if (upperIndCount == lowerIndCount)
   return pointer;

  return -1;
 }

If this is the optimum or not im not sure. Im also not sure if this will work on very large data sets. But for the data sets i tried(both negative and positive) it worked fine. And also this has O(N) complexity as I am iterating through the array only once.