Three Equal Parts

Intuition

Solution

class Solution:
    def threeEqualParts(self, A: List[int]) -> List[int]:
        i = 0
        tmp = 0
        while i < len(A) and A[i] == 0:
            i += 1
            tmp += 1
        A = A[i:]
        if not A:
            return [0, 2]
        elif len(A) < 3:
            return [-1, -1]

        zeroes = [0] * len(A)
        if not A[-1]:
            zeroes[-1] = 1
        for i in range(len(A) - 2, -1, -1):
            if not A[i]:
                zeroes[i] = zeroes[i + 1] + 1

        for l in range(1, len(A) // 2 + 1):
            if zeroes[l] + l == len(A):
                return [-1, -1]
            if zeroes[l] + l + l < len(A) and 3 * l + zeroes[l] + zeroes[zeroes[l] + l + l] >= len(A):
                if self.match(A, l, 2 * l + zeroes[l]):
                    return [l - 1 + tmp, 2 * l + zeroes[l] + tmp]
        return [-1, -1]

    def match(self, A, i, j):
        s1, s2, s3 = 0, i, j
        while s2 < len(A) and A[s2] == 0:
            s2 += 1
        while s3 < len(A) and A[s3] == 0:
            s3 += 1

        if s3 == len(A) or s2 == len(A):
            return False
        while s3 < len(A):
            if A[s1] == A[s2] == A[s3]:
                s1 += 1
                s2 += 1
                s3 += 1
            else:
                return False
        if s2 == j and s3 == len(A):
            return True
        return False