927. Three Equal Parts

You are given an array arr which consists of only zeros and ones, divide the array into three 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:

• arr[0], arr[1], …, arr[i] is the first part,
  varr[i + 1], arr[i + 2], …, arr[j - 1] is the second part, and
• arr[j], arr[j + 1], …, arr[arr.length - 1] is the third part.
• All three parts have equal binary values.

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.

class Solution {
    bool match(int l, int r, int end, vector<int>& A) {
        int len1 = r - l, len2 = A.size() - 1 - end;
        if(len1 < len2) return false;
        int p = A.size() - 1;
        while(p >= end) if(A[r--] != A[p--]) return false;
        while(r >= l) if(A[r--] == 1) return false;
        return true;
    }
    vector<int> zfunction(vector<int>& now) {
        int l = 0, r = 0, n = now.size();
        vector<int> z(n);
        for(int i = 1; i < n; i++) {
            if(i > r) {
                l = r = i;
                while(r < n and now[r-l] == now[r]) r++;
                z[i] = r - l;
                r--;   
            } else {
                int k = i - l;
                if(z[k] < r - i + 1) z[i] = z[k];
                else {
                    l = i;
                    while(r < n and now[r-l] == now[r]) r++;
                    z[i] = r - l;
                    r--;
                }
            }
        }
        return z;
    }
public:
    vector<int> threeEqualParts(vector<int>& A) {
        int cnt = count(begin(A), end(A), 1);
        if(cnt == 0) return {0,(int)A.size() - 1};
        if(cnt % 3) return {-1,-1};
        
        int cut = 0;
        for(int i = 0; i < A.size() and A[i] == 0; i++) cut += 1;
        
        vector<int> now;
        for(int i = cut; i < A.size(); i++) now.push_back(A[i]);
        
        int n = now.size();
        vector<int> z = zfunction(now);
        
        int padding = 0;
        for(int i = n - 1; now[i] == 0; i--) padding++;
        
        for(int i = n - 1, req = 1; i >= 0; i--,req++) {
            if(z[i] != req) continue;
            
            int j = i - 1;
            while(j >= req and now[j] == 0) j--;
            j += padding;
            
            if(j >= i) continue;
            if(j == req - 1) return {-1,-1};
            if(match(req,j,i,now)) return {cut + req - 1, cut + j + 1};
            
            req += (i - j - 1);
            i = j + 1;
            
        }
        return {-1,-1};
    }
};