Leetcode 题目927-三等分0和1组成的数组

class Solution:

    def threeEqualParts(self, arr: list[int]) -> list[int]:
        n = len(arr)
        if n < 3: return [-1, -1]
        s = sum(arr)  # arr 中 1的个数
        if s % 3 != 0: return [-1, -1]  # 1 的个数不能被3整除则无解
        if s == 0 and len(arr) >= 3: return [0, 2]

        m = s//3  # qty of ones in each
        flag = True
        s1 = 0
        for i in range(n):
            if arr[i] != 0:
                s1 += 1
                if flag:
                    ones1_left = i
                    flag = False
                if s1 == m:
                    ones1_right = i
                    break

        flag = True
        s2 = 0
        for j in range(i+1, n):
            if arr[j] != 0:
                s2 += 1
                if flag:
                    ones2_left = j
                    flag = False
                if s2 == m:
                    ones2_right = j
                    break


        flag = True
        s3 = 0
        for k in range(j+1, n):
            if arr[k] != 0:
                s3 += 1
                if flag:
                    ones3_left = k
                    flag = False
                if s3 == m:
                    ones3_right = k
                    break


        right_side_zeros = n - (ones3_right + 1)  # 最后一组末尾的0的个数
        # print(right_side_zeros)
        if ones2_left < ones1_right + right_side_zeros + 1 or ones3_left < ones2_right + right_side_zeros + 1:
            return [-1, -1]


        i = ones1_right + right_side_zeros
        a1 = arr[ones1_left: i+1]
        if n < 3 * len(a1):
            return [-1,-1]

        j = ones2_right + right_side_zeros + 1

        if (j-i-1) < len(a1) or (n-j) < len(a1):  # a2 或 a3 的可能的最大长度不能小于a1的长度
            return [-1,-1]
        a2 = arr[ones2_left: j]
        if a1 != a2:
            return [-1,-1]

        a3 = arr[ones3_left:]
        if a1 == a3:
            return [i,j]

        return [-1,-1]