1. 程式人生 > >LeetCode 927. 三等分

LeetCode 927. 三等分

給定一個由 0 和 1 組成的陣列 A,將陣列分成 3 個非空的部分,使得所有這些部分表示相同的二進位制值。

如果可以做到,請返回任何 [i, j],其中 i+1 < j,這樣一來:

  • A[0], A[1], ..., A[i] 組成第一部分;
  • A[i+1], A[i+2], ..., A[j-1] 作為第二部分;
  • A[j], A[j+1], ..., A[A.length - 1] 是第三部分。
  • 這三個部分所表示的二進位制值相等。

如果無法做到,就返回 [-1, -1]

注意,在考慮每個部分所表示的二進位制時,應當將其看作一個整體。例如,[1,1,0] 表示十進位制中的 6,而不會是 3

。此外,前導零也是被允許的,所以 [0,1,1] 和 [1,1] 表示相同的值。

示例 1:
輸入:[1,0,1,0,1]
輸出:[0,3]

示例 2:
輸出:[1,1,0,1,1]
輸出:[-1,-1]

提示:
3 <= A.length <= 30000
A[i] == 0 或 A[i] == 1

演算法 

就是把陣列分成3部分,先保證每部分的1的個數是一樣的,以及尾部的0的個數,再轉換成整數形式判斷值是否相等。

程式碼 

class Solution:
    def threeEqualParts(self, A):
        """
        :type A: List[int]
        :rtype: List[int]
        """
        sumA = sum(A)
        if sumA == 0:
            return [0, 2]
        if sumA % 3:
            return [-1, -1]
        cal = sumA // 3
        tail0, i = 0, len(A) - 1
        while A[i] == 0:
            i -= 1
            tail0 += 1
        lst = [[], [], []]
        tmp, idx, tail = 0, 0, 0
        for a in A:
            if tmp < cal:
                tmp += a
                lst[idx].append(str(a))
            elif tail < tail0:
                if a == 0:
                    lst[idx].append(str(a))
                    tail += 1
                else:
                    return [-1, -1]
            else:
                tmp, tail = 0, 0
                idx += 1
                tmp += a
                lst[idx].append(str(a))
        if int("".join(lst[0])) == int("".join(lst[1])) == int("".join(lst[2])):
            return [len(lst[0]) - 1, len(lst[0]) + len(lst[1])]
        else:
            return [-1, -1]