Given an array A of integers, return true if and only if we can partition the array into three non-empty parts with equal sums.

Formally, we can partition the array if we can find indexes i+1 < j with (A[0] + A[1] + … + A[i] == A[i+1] + A[i+2] + … + A[j-1] == A[j] + A[j-1] + … + A[A.length – 1])

Example 1:
Input: [0,2,1,-6,6,-7,9,1,2,0,1]
Output: true
Explanation: 0 + 2 + 1 = -6 + 6 – 7 + 9 + 1 = 2 + 0 + 1

Example 2:
Input: [0,2,1,-6,6,7,9,-1,2,0,1]
Output: false

Example 3:
Input: [3,3,6,5,-2,2,5,1,-9,4]
Output: true
Explanation: 3 + 3 = 6 = 5 – 2 + 2 + 5 + 1 – 9 + 4

Note:
3 <= A.length <= 50000
-10000 <= A[i] <= 10000

Array Paritition Algorithm via Brute-force Algorithm

We can have two index pointers i, and j, which runs at O(N^2) time complexity. We also keep updated partial sum from 0 to i and i to j respectively, then we need to check if these two paritial sums are equal and also equal to the remainder.

This brute force algorithm is slow, which gives a time limit exceeded error if the input list is huge.

Compute the Average

We can do a O(N) to compute the sum first. Then, we know the one third of the sum. When we go through the array, we accumulate the sum, if it is equal to the average, we increment the counter and reset the sum.

At the end of the array O(N), if the counter is three and the current sum is zero, we know that the list can be divided into perfectly 3 equal parts.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

class Solution {
public:
    bool canThreePartsEqualSum(vector<int>& A) {
        int avg = std::accumulate(begin(A), end(A), 0, [](auto &a, auto &b) {return a + b; }) / 3;
        int cur = 0;
        int count = 0;
        for (int i = 0; i < A.size(); ++ i) {
            cur += A[i];
            if (cur == avg) {
                count ++;
                cur = 0;
            }
        }
        return count == 3 && cur == 0;
    }
};

class Solution {
public:
    bool canThreePartsEqualSum(vector<int>& A) {
        int avg = std::accumulate(begin(A), end(A), 0, [](auto &a, auto &b) {return a + b; }) / 3;
        int cur = 0;
        int count = 0;
        for (int i = 0; i < A.size(); ++ i) {
            cur += A[i];
            if (cur == avg) {
                count ++;
                cur = 0;
            }
        }
        return count == 3 && cur == 0;
    }
};

To compute the sum, we can use the std::accumulate instead of the tradition for-loop.

–EOF (The Ultimate Computing & Technology Blog) —