927. Three Equal Parts

Given an array A of 0s and 1s, divide the array into 3 non-empty parts such that all of these parts represent the same binary value.

If it is possible, return any [i, j] with i+1 < j, such that:

• A[0], A[1], ..., A[i] is the first part;
• A[i+1], A[i+2], ..., A[j-1] is the second part, and
• A[j], A[j+1], ..., A[A.length - 1] is the third part.
• All three parts have equal binary value.

If it is not possible, return [-1, -1].

Note that the entire part is used when considering what binary value it represents.  For example, [1,1,0] represents 6 in decimal, not 3.  Also, leading zeros are allowed, so [0,1,1] and [1,1] represent the same value.

 

Example 1:

Input: [1,0,1,0,1]
Output: [0,3]

Example 2:

Input: [1,1,0,1,1]
Output: [-1,-1]

 

Note:

1. 3 <= A.length <= 30000
2. A[i] == 0 or A[i] == 1

 

class Solution {
public:
    vector<int> threeEqualParts(vector<int>& A) {
        int size = A.size();
        int countOfOne = 0;
        for (auto c : A)
            if (c == 1)
                countOfOne++;
        
        // if there don't have 1 in the vector
        if (countOfOne == 0)
            return {0, size-1};
        
        // if the count of one is not a multiple of 3, then we can never find a possible partition since
        //there will be at least one partion that will have difference number of one hence different binary
        //representation
        //For example, given:
        //0000110   110   110
        //    |     |     |
        //    i           j
        //Total number of ones = 6
        if (countOfOne%3 != 0)
            return {-1, -1};
        int k = countOfOne / 3;
        int i;
        for (i = 0; i < size; ++i) 
            if (A[i] == 1)
                break;
        int begin = i;
        int temp = 0;
        for (i = 0; i < size; ++i) {
            if (A[i] == 1)
                temp++;
            if (temp == k + 1)
                break;
        }
        int mid = i;
        temp = 0;
        for (i = 0; i < size; ++i) {
            if (A[i] == 1)
                temp++;
            if (temp == 2*k+1)
                break;
        }
        int end = i;
        while (end < size && A[begin] == A[mid] && A[mid] == A[end]) {
            begin++, mid++, end++;
        }
        if (end == size)
            return {begin-1, mid};
        else 
            return {-1, -1};
    }
};