Drop two elements to split the array to three part evenly in O(n)

Question

I encounter a problem to let you drop two elements in an array to make the three part's sum equal.

  Ex:
  1 2 4 3 5 2 1
  After I drop the 4 and 5, it becomes 1, 2 | 3 | 2, 1

Constraints:

  1.Numbers are all integer > 0

  2.Drop two elements in the array, so the three splitted subarrays will have same subarray sum.

I have tried it by using two pass algorithm as the following

First pass:O(n) Count the accumulated sum from the left so I can get the range sum easily.

Second pass:O(n^2) Use nested loop to get the subarray's start and end index. Then, calculate the left, mid, right sum.

// 1.get accumulated sum map
int[] sumMap = new int[A.length];
int sum = 0;
for(int i = 0; i < A.length; i ++) {
    sum += A[i];
    sumMap[i] = sum;
}

// 2.try each combination
for(int i = 1; i < A.length - 1; i ++) {
    for(int j = i + 1; j < A.length - 1; j ++) {
        int left = sumMap[i] - A[i];
        int mid = sumMap[j] - sumMap[i] - A[j];
        int right = sumMap[A.length - 1] - sumMap[j];

        if(left == mid && mid == right)return true;
    }
}

Are there any better algorithm that can achieve O(n)?

Thanks

Answer 1

Assuming that first and last element can't be dropped and all elements are >0:

Set a variable sumleft to value of first element, sumright to value of last element. You also need index variables to remember which elements from left and right were already added to the sums.

1. If sumleft == sumright, test if next elements from left and right can be dropped to fulfill requirement. If so -> done. If not take next elements from left and right and add it to the respective sum variable. Back to 1.

2. If sumleft < sumright, add next value from the left to sumleft. Back to 1.

3. If sumleft > sumright, add next value from the right to sumright. Back to 1.

If all elements were consumed, there is no solution.

Edit: Testing if requirement is fulfilled when sumleft == sumright can be done by initially summing up all elements (also needs only O(n)) and checking if this sum minus the elements to drop is equal sumleft * 3.