leetcode 927.threeEqualParts 三等分

leetcode 927.threeEqualParts

class Solution:
    def threeEqualParts(self, A):
        """
        :type A: List[int]
        :rtype: List[int]
        """
        A = [str(a) for a in A]
        for i in range(0, len(A) - 2):            
            for j in range(i + 2, len(A)):                
                if int(str(A[0]) + ''.join(A[1:i+1]), 2) == int(''.join(A[i + 1:j]), 2) == int(''.join(A[j:]), 2):
                    return [i, j]
        else:
            return [-1, -1]
a = Solution()
l=[0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,1,1,1,1,0,1,1,1,1,1,0,0,0,0,1,0,0,1,1,0,0,0,1,1,1,0,1,0,1,1,1,1,0,1,1,0,1,1,1,1,0,1,1,0,0,0,0,0,0,0,0,1,0,0,0,0,1,0,0,0,0,0,0,0,0,1,0,1,1,1,1,1,1,1,0,0,0,0,0,0,0,1,0,1,1,1,1,1,0,0,1,1,0,1,0,1,0,0,1,1,0,1,1,0,1,0,0,0,1,1,0,0,0,1,1,0,0,1,0,1,0,1,1,0,0,0,1,0,0,0,0,1,1,1,1,0,0,1,1,0,0,0,0,1,1,0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,1,1,1,1,0,1,1,1,1,1,0,0,0,0,1,0,0,1,1,0,0,0,1,1,1,0,1,0,1,1,1,1,0,1,1,0,1,1,1,1,0,1,1,0,0,0,0,0,0,0,0,1,0,0,0,0,1,0,0,0,0,0,0,0,0,1,0,1,1,1,1,1,1,1,0,0,0,0,0,0,0,1,0,1,1,1,1,1,0,0,1,1,0,1,0,1,0,0,1,1,0,1,1,0,1,0,0,0,1,1,0,0,0,1,1,0,0,1,0,1,0,1,1,0,0,0,1,0,0,0,0,1,1,1,1,0,0,1,1,0,0,0,0,1,1,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,1,1,1,1,0,1,1,1,1,1,0,0,0,0,1,0,0,1,1,0,0,0,1,1,1,0,1,0,1,1,1,1,0,1,1,0,1,1,1,1,0,1,1,0,0,0,0,0,0,0,0,1,0,0,0,0,1,0,0,0,0,0,0,0,0,1,0,1,1,1,1,1,1,1,0,0,0,0,0,0,0,1,0,1,1,1,1,1,0,0,1,1,0,1,0,1,0,0,1,1,0,1,1,0,1,0,0,0,1,1,0,0,0,1,1,0,0,1,0,1,0,1,1,0,0,0,1,0,0,0,0,1,1,1,1,0,0,1,1,0,0,0,0,1,1,0,0,0,1,0,0,0,0,0,0,0]

timeit

301 ms ± 3.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

提交以后超时

测试

• any(l): 171 ns ± 0.689 ns
• all(l): 95.5 ns ± 8.28 ns
  优化代码

class Solution:
    def threeEqualParts(self, A):
        """
        :type A: List[int]
        :rtype: List[int]
        """
        if all([i == 0 for i in A]):
            return [0, 2]
        else:
            A = [str(a) for a in A]
            i = 0
            j = len(A) - 1
            while i + 1 < j:
                if int(str(A[0]) + ''.join(A[1:i + 1]), 2) > int(''.join(A[j:]), 2):
                    j -= 1
                elif int(str(A[0]) + ''.join(A[1:i + 1]), 2) < int(''.join(A[j:]), 2):
                    i += 1
                elif int(str(A[0]) + ''.join(A[1:i + 1]), 2) == int(''.join(A[j:]), 2):
                    if int(str(A[0]) + ''.join(A[1:i + 1]), 2) == int(''.join(A[i + 1:j]), 2):
                        return [i, j]
                    else:
                        j -= 1
            else:
                return [-1, -1]

速度有所提高：1.88 ms ± 68.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

提交仍然超时

看答案：

class Solution:
    def threeEqualParts(self, A):
        """
        :type A: List[int]
        :rtype: List[int]
        """
        # 如果全零返回[0, 2]
        if all([i == 0 for i in A]):
            return [0, 2]
        
        # 要满足条件，列表中1的数量必然可以被3整除，t是分成3段每段中1的个数
        if sum(A) % 3 != 0:
            return [-1, -1]
        else:
            t = sum(A) / 3
        
        # 列表切分成3段，m1-m6是每段的第一个1和最后一个1的位置，分两步追加是因为m1与m4、m2与m5、m3与m6可能重合
        interval = []
        flag = 0
        for i, n in enumerate(A):
            if n:
                flag += 1
                if flag in (1, t + 1, 2 * t + 1):
                    interval.append(i)
                if flag in (t, 2 * t, 3 * t):
                    interval.append(i)
        m1, m2, m3, m4, m5, m6 = interval
        
        # 要想3段等分，首先要满足3段中首尾是1的部分完全相同
        if A[m1: m2 + 1] == A[m3: m4 + 1] == A[m5: m6 + 1]:
            # 然后第一段和第二段的最后一个1后面的0的个数要大于等于最后一段最后一个1后面的0的个数
            zero_after_m6 = len(A) - (m6 + 1)
            if len(A[m2+1:m3]) >= zero_after_m6 and len(A[m4+1:m5]) >= zero_after_m6:
                return [m2 + zero_after_m6, m4 + zero_after_m6 + 1]
            else:
                return [-1, -1]
        else:
            return [-1, -1]

本地测试216 µs ± 31.5 ，提交通过