[LeetCode]1013.将数组分成和相等的三个部分

题目

给你一个整数数组 A，只有可以将其划分为三个和相等的非空部分时才返回 true，否则返回 false。

形式上，如果可以找出索引 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：

输出：[0,2,1,-6,6,-7,9,1,2,0,1]
输出：true
解释：0 + 2 + 1 = -6 + 6 - 7 + 9 + 1 = 2 + 0 + 1

示例 2：

输入：[0,2,1,-6,6,7,9,-1,2,0,1]
输出：false

示例 3：

输入：[3,3,6,5,-2,2,5,1,-9,4]
输出：true
解释：3 + 3 = 6 = 5 - 2 + 2 + 5 + 1 - 9 + 4

提示：

1. 3 <= A.length <= 50000
2. -10^4 <= A[i] <= 10^4

解题思路

题目要求数组能被分成和相等的三个部分，那么数组之和必然是3的倍数，如果不满足这个先决条件，则可以直接返回False。
在满足数组和是3的倍数之后，就从前往后遍历数组，则每一个非空部分的和都应当是 sum(A) / 3。因此我们需要找到索引 i 和 j 使得：

A[0] + A[1] + … + A[i] = sum(A) / 3;

A[i + 1] + A[i + 2] + … + A[j] = sum(A) / 3。这等价于 A[0] + A[1] + … + A[j] = sum(A) / 3 * 2 且 j > i。

首先需要找出索引 i。具体地，我们从第一个元素开始遍历数组 A 并对数组中的数进行累加。当累加的和等于 sum(A) / 3 时，我们就将当前的位置置为索引 i。
我们从 i + 1 开始继续遍历数组 A 并进行累加，当累加的和等于 sum(A) / 3 * 2 时，我们就得到了索引 j，可以返回 true 作为答案。如果我们无法找到索引 i 或索引 j，那么返回 false。

复杂度分析：
时间复杂度：O(N)，其中 N 是数组 A 的长度。我们最多只需要遍历一遍数组就可以得到答案。
空间复杂度：O(1)。我们只需要使用额外的索引变量 i，j 以及一些存储数组信息的变量。

代码

class Solution:
    def canThreePartsEqualSum(self, A: List[int]) -> bool:
        s = sum(A)
        if s%3 != 0:
            return False
        else:
            r = int(s//3)
            n = len(A)
            i, cur = 0, 0
            while i < n:
                cur += A[i]
                if cur == r:
                    break
                i += 1
            if cur != r:
                return False
            j = i+1
            while j+1 < n: # 需要满足最后一个数组非空
                cur += A[j]
                if cur == 2*r:
                    return True
                j += 1
            return False