Finding an equilibrium index in an int array

Question

Here is a programming challenge from codility

A zero-indexed array A consisting of N integers is given. An equilibrium index of this array is any integer P such that 0 ≤ P < N and the sum of elements of lower indices is equal to the sum of elements of higher indices, i.e.

A[0] + A[1] + ... + A[P−1] = A[P+1] + ... + A[N−2] + A[N−1].

Sum of zero elements is assumed to be equal to 0. This can happen if P = 0 or if P = N−1.

For example, consider the following array A consisting of N = 8 elements:

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

P = 1 is an equilibrium index of this array, because:

• A[0] = −1 = A[2] + A[3] + A[4] + A[5] + A[6] + A[7]

P = 3 is an equilibrium index of this array, because:

• A[0] + A[1] + A[2] = −2 = A[4] + A[5] + A[6] + A[7]

P = 7 is also an equilibrium index, because:

• A[0] + A[1] + A[2] + A[3] + A[4] + A[5] + A[6] = 0

and there are no elements with indices greater than 7.

P = 8 is not an equilibrium index, because it does not fulfill the condition 0 ≤ P < N.

Write a function:

class Solution { public int solution(int[] A); }

that, given a zero-indexed array A consisting of N integers, returns any of its equilibrium indices. The function should return −1 if no equilibrium index exists.

For example, given array A shown above, the function may return 1, 3 or 7, as explained above.

Assume that:

• N is an integer within the range [0..100,000];
• each element of array A is an integer within the range [−2,147,483,648..2,147,483,647].

Complexity:

• expected worst-case time complexity is \$O(N)\$;
• expected worst-case space complexity is \$O(N)\$, beyond input storage (not counting the storage required for input arguments).

Elements of input arrays can be modified.

• I think my whole solution is \$O(N)\$ because both my methods are \$O(N)\$. Is that right?
• I believe I used some space to generate the sums and it Is still \$O(N)\$. Is it right?

Please review my code:

class Equilibrium {

    public int getEquilibrium(int[] A) {
        long[] sums = generateSums(A);
        long lowSum = 0L;
        int res = -1;
        for (int i = 0; i < A.length; i++) {
            if (lowSum == sums[i+1]) {
                res = i;
                break;
            }
            lowSum += A[i];
        }
        return res;
    }

    // I used long to handle sums greater than 32 bits
    public long[] generateSums(int[] A) {
        // I added another index (default value is 0) to handle the last value (E.G. the sum of previous elements is 0, so the equilibrium index should be the last value because there are no other items after it. I used the newly added index to handle it.)
        long[] res = new long[A.length+1];
        for (int i = A.length-1; i >= 0; i--) {
            res[i] = i+1 == A.length ? A[i] : res[i+1]+A[i];
        }
        return res;
    }
}

Did I follow all the requirements in the problem? How can I improve this and make it faster?

Answer 1

• I think my whole solution is \$O(N)\$ because both my methods are O(N). Is that right?

Correct. Both of your methods are \$O(N)\$: they are linear in terms of the number of elements in the input array. More specifically, both generateSums and getEquilibrium loop over the array once.

• I believe I used some space to generate the sums and it Is still \$O(N)\$. Is it right?

Also correct. The intermediate array containing the sum contains as many elements as there was in the input array; doubling the input array in size would double the size of the intermediate array. This makes in linear in terms of space complexity.

---

The method generateSums doesn't have a good name. Its purpose is to calculate the cumulative sum of the input array, going from the right to the left. Consider naming it cumulativeSums instead. The rest of the algorithm is correct. You won't be able to make this faster than \$O(N)\$, in terms of time complexity.

A small comment: instead of using a break explicitly inside the for loop, consider returning directly the found index. You currenly have:

int res = -1;
for (int i = 0; i < A.length; i++) {
    if (lowSum == sums[i+1]) {
        res = i;
        break;
    }
    lowSum += A[i];
}
return res;

but it would be shorter to have:

for (int i = 0; i < A.length; i++) {
    if (lowSum == sums[i+1]) {
        return i;
    }
    lowSum += A[i];
}
return -1;

It avoids the use of a temporary local variable holding the result, and makes it clearer that -1 is the default value returned.

In terms of space complexity however, you can do better: this is possible in \$O(1)\$. Instead of generating an array containing the cumulative sums of the array, just calculate the total sum, let's call it right sum. Then, you keep a running sum of the array by traversing it from left to right, decreasing this right sum as you go: if the running sum is equal to the right sum, you found an equilibrium.

public int getEquilibrium(int[] array) {
    long totalSum = sum(array);
    long lowSum = 0L;
    for (int i = 0; i < array.length; i++) {
        totalSum -= array[i];
        if (lowSum == totalSum) {
            return i;
        }
        lowSum += array[i];
    }
    return -1;
}

public long sum(int[] array) {
    return Arrays.stream(array).asLongStream().sum();
}

Answer 2

I think my whole solution is \$O(N)\$ because both my methods are \$O(N)\$. Is that right?

That's correct. You traverse the \$N\$ values twice, the time complexity is \$O(N)\$.

I believe I used some space to generate the sums and it Is still \$O(N)\$. Is it right?

Yes, you used extra \$O(N)\$ space.

Did I follow all the requirements in the problem? How can I improve this and make it faster?

You followed the requirements, and apparently you got perfect score. But you can do better, by using \$O(1)\$ instead of \$O(N)\$.

You don't need to store the sums in an array. You can calculate the prefix sum as you go, and return the index when prefix == sum - prefix + A[i]. This also implies that you don't need int res = -1;, you can return the index immediately when you find it, and return -1 if the loop terminates without finding an equilibrium index.

public int findEquilibriumIndex(int[] arr) {
    long sum = sum(arr);

    long prefix = 0;
    for (int index = 0; index < arr.length; index++) {
        prefix += arr[index];
        if (prefix == sum - prefix + arr[index]) {
            return index;
        }
    }
    return -1;
}

public long sum(int[] arr) {
    long sum = 0;
    for (int value : arr) {
        sum += value;
    }
    return sum;
}