leetcode1712.将数组分成三个子数组的方案数

1.题目描述

1712. 将数组分成三个子数组的方案数

我们称一个分割整数数组的方案是 好的 ，当它满足：

• 数组被分成三个 非空 连续子数组，从左至右分别命名为 left ， mid ， right 。
• left 中元素和小于等于 mid 中元素和，mid 中元素和小于等于 right 中元素和。

给你一个 非负 整数数组 nums ，请你返回 好的 分割 nums 方案数目。由于答案可能会很大，请你将结果对 109 + 7 取余后返回。

示例 1：

输入：nums = [1,1,1]
输出：1
解释：唯一一种好的分割方案是将 nums 分成 [1] [1] [1] 。

示例 2：

输入：nums = [1,2,2,2,5,0]
输出：3
解释：nums 总共有 3 种好的分割方案：
[1] [2] [2,2,5,0]
[1] [2,2] [2,5,0]
[1,2] [2,2] [5,0]

示例 3：

输入：nums = [3,2,1]
输出：0
解释：没有好的分割方案。

提示：

• 3 <= nums.length <= 10**5
• 0 <= nums[i] <= 10**4

2.解题思路

这道题断断续续思考了有两天

先看初始思路：

固定第一刀 i 的分割位置，遍历第二刀

第一刀的起始位置：可以从 0 开始。
第一刀的终止位置：最大到 sum(nums) / 3位置（可以等于），因为超过该位置，则不能满足sum(left) <= sum(mid) <= sum(right)

在遍历第二刀 j 过程中，

需要满足以下条件

sum[mid] >= sum[left]

sum[mid] <= sum[right]---> sum[mid] + sum[right] <= 2 * sum[right]

为了能在O(1)时间快速计算每一段的和，使用前缀和数组sums

因此，上述两个条件转化成：

不等式1：sums[j] - sums[i] >= sums[i]

不等式2：sums[n-1] - sums[i] <= 2 * (sums[n-1] - sums[j])

因此有了如下代码，10**5的数组，两次循环遍历，超时

import itertools

class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        n = len(nums)
        sums = list(itertools.accumulate(nums))
        MOD = 1000000000 + 7
        ans = 0
        for i in range(0,len(nums)-1):
            if sums[i] <= sums[n - 1] / 3:
                for j in range(i+1,len(nums)-1):
                    if sums[j] >= 2 * sums[i] and sums[n-1] - sums[i] <= 2 * (sums[n-1] - sums[j]):
                        ans +=  1
        return ans % MOD

优化思路：

将内层循环优化为二分查找

根据不等式1，可以推出第二刀j的起始位置要满足 sums[j] >= 2 * sums[i]，因此利用二分查找中的 lower_bound思路，找到在前缀和数组，大于等于目标值2 * sums[i]的第一个索引

根据不等式2，可以推出第二刀j的结束位置要满足 sums[j] <= (sums[n-1] + sums[i]) / 2，因此利用二分查找中的 upper_bound思路，找到在前缀和数组，大于目标值(sums[n-1] + sums[i]) / 2的第一个索引

注意，如果这里的sums[j]是大于等于目标值的最后一个索引，其实是sums[j] == (sums[n-1] + sums[i]) / 2的情况

注意这里为什么要使用upper_bound

 于是有了下面代码，时间复杂度还是过不去！

import itertools
import bisect
class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        n = len(nums)
        sums = list(itertools.accumulate(nums))
        MOD = 1000000000 + 7
        ans = 0
        for i in range(n-1):
            if sums[i] > sums[n-1] / 3:
                break
            j1 = bisect.bisect_left(sums[i+1:], 2 * sums[i])
            j2 = bisect.bisect_right(sums[i+1:], (sums[n-1] + sums[i]) / 2)
            if j2 >= j1:
                ans += (j2-j1)
        return ans % MOD

sums[i+1:]这一行时间复杂度O(n)的，改成直接在前缀和数组上二分查找就行，最终代码实现如下：

3.代码实现

import itertools
import bisect
class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        n = len(nums)
        sums = list(itertools.accumulate(nums))
        MOD = 1000000000 + 7
        ans = 0
        # 从零开始，代表分割了第一个数
        for i in range(n-1):
            if sums[i] > sums[n-1] / 3:
                break
            j1 = bisect.bisect_left(sums, 2 * sums[i])
            j1 = max(i+1,j1)
            j2 = bisect.bisect_right(sums, (sums[n-1] + sums[i]) / 2)
            # j2是大于目标值的第一个数，取不到，所以min中取n-1
            j2 = min(n-1,j2)
            if j2 >= j1:
                ans += (j2-j1)
        return ans % MOD