LeetCode每日一题——927. 三等分

LeetCode每日一题系列

题目：927. 三等分
难度：困难

---

---

题目

给定一个由 0 和 1 组成的数组 arr ，将数组分成 3 个非空的部分 ，使得所有这些部分表示相同的二进制值。

如果可以做到，请返回任何 [i, j]，其中 i+1 < j，这样一来：

arr[0], arr[1], …, arr[i] 为第一部分；
arr[i + 1], arr[i + 2], …, arr[j - 1] 为第二部分；
arr[j], arr[j + 1], …, arr[arr.length - 1] 为第三部分。
这三个部分所表示的二进制值相等。
如果无法做到，就返回 [-1, -1]。

注意，在考虑每个部分所表示的二进制时，应当将其看作一个整体。例如，[1,1,0] 表示十进制中的 6，而不会是 3。此外，前导零也是被允许的，所以 [0,1,1] 和 [1,1] 表示相同的值。

示例

示例 1：

输入：arr = [1,0,1,0,1]
输出：[0,3]

示例 2：

输入：arr = [1,1,0,1,1]
输出：[-1,-1]

示例 3:

输入：arr = [1,1,0,0,1]
输出：[0,2]

提示：

3 <= arr.length <= 3 * 104
arr[i] 是 0 或 1

思路

自己写的双重for循环暴力寻找分割方法很明显会超时，借鉴了官方的做法

1. 首先统计1的个数，如果1的个数不是3的倍数，则不可能三等分，直接返回即可。如果1的个数为零，这时可以返回符合题意的答案[0,2]
2. 每一部分都应该有（index = len(arr) // 3）个1，我们需要找到第一个1、第index+1个1、第2*index+1个1的下标。因为给定数组的长度是确定的，所以我们的第三个部分的长度是已经确定的，也就是（n = len(arr) - third）
3. 分别比较从first、second、third开始的之后n位是否全部相等（在下标否符合条件的情况下，也就是first + l <= second and second + l <= third），如果全部相等可以返回答案[first+l-1, second+l]，如果不相等或下标不符合条件，直接返回[-1,-1]

题解

class Solution:
    def threeEqualParts(self, arr: List[int]) -> List[int]:
        # 计算1的数量
        s = 0
        for i in range(len(arr)):
            if arr[i] == 1:
                s += 1
        # 1的数量不为3的倍数，直接返回
        if s % 3 != 0:
            return [-1,-1]
        # 没有1，直接返回
        if s == 0:
            return [0,2]
        index = s // 3
        # 寻找第一个、第index个、第2*index个1
        first, second, third, cur = 0,0,0,0
        for i,j in enumerate(arr):
            if j == 1:
                if cur == 0:
                    first = i
                if cur == index:
                    second = i
                if cur == 2 * index:
                    third = i
                cur += 1
        # 比较从first，second，third相对位置是否全部相同
        l = len(arr) - third
        # 下标需符合条件
        if first + l <= second and second + l <= third:
            n = 1
            while third + n <len(arr):
                # 不相等直接返回
                if arr[first + n] != arr[second + n] or arr[first + n] != arr[third + n]:
                    return [-1,-1]
                n += 1
            # 最后全部相等返回
            return [first+l-1, second+l]
        return [-1,-1]