LeetCode 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$

由于输入数组nums中的值都大于等于0，因此其前缀和数组是非递减的。我们可以枚举中间子数组的右端点rm，对于一个右端点rm，为了使中间子数组的和sm大于等于左边子数组的和sl，我们可以找到一个最大的中间子数组的左端点lm1，显然lm1左移都可以使sm>=sl；为了使右边子数组的和sr大于等于中间子数组的和sm，我们可以找到一个最小的中间子数组的左端点lm2，显然lm2右移都可以使sr>=sm；此时lm1和lm2之间的元素即为合法的lm。由于nums的前缀和数组具有单调性，因此rm右移时，lm1和lm2不会左移：

class Solution {
public:int waysToSplit(vector<int>& nums) {vector<int> preSum(nums.size() + 1);for (int i = 0; i < nums.size(); ++i) {preSum[i + 1] = preSum[i] + nums[i];}long long ways = 0;int lowerLeft = 1;int upperLeft = 1;// 从1开始枚举，到nums.size() - 2，保证了左右子数组不为空for (int i = 1; i < nums.size() - 1; ++i) {// 注释里的代码更易理解，即找到使sm>=sl的最右点再右移一位的位置// while (upperLeft <= i && preSum[upperLeft] <= preSum[i + 1] - preSum[upperLeft]) {while (upperLeft <= i && preSum[upperLeft] * 2 <= preSum[i + 1]) {++upperLeft;}// 注释里的代码更易理解，即找到使sm<=sl的最左点// while (lowerLeft <= i && preSum[i + 1] - preSum[lowerLeft] > preSum[nums.size()] - preSum[i + 1]) {while (lowerLeft <= i && 2 * preSum[i + 1] > preSum[nums.size()] + preSum[lowerLeft]) {++lowerLeft;}// [lowerLeft, upperLeft)之间的取值即为lm的合法取值ways += max(upperLeft - lowerLeft, 0);}return ways % (long long)(1e9 + 7);}
};

如果nums的长度为n，则此算法时间复杂度为O(n)，外层for循环n次，内层两个while最多只会各执行n次，均摊到每次for循环都是O(1)；空间复杂度为O(n)，前缀和数组的长度为n+1。