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

题目

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

---

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

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

给你一个 非负 整数数组 nums ，请你返回 好的 分割 nums 方案数目。由于答案可能会很大，请你将结果对 10^9 + 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[i] i from 0 to n-1

题目要求nums[0...i] <= nums[i+1...j] <= nums[j+1...n-1]，其中i<j, i from 0 to n-3, j from 1 to n-2。

令p[i] = nums[0...i-1]，得p[i+1]-p[0]<=p[j+1]-p[i+1]<=p[n]-p[j+1]
变型可得$2p[i+1] \le p[j+1] \le \lfloor \frac{p[n]+p[i+1]}{2} \rfloor$

变量替换得

$2p[u] \le p[v] \le \lfloor \frac{p[n]+p[u]}{2} \rfloor, u<v, u \in [1, n-2], v \in [2, n-1]$

由此我们可以遍历u，符合条件的v的个数是非负的r-l+1

其中l是p中第一个大于等于2p[u]的下标与下标i+1的最大值。

其中r是p中最后一个小于等于$\lfloor \frac{p[n]+p[u]}{2} \rfloor$的下标与n-1的最小值。

代码

class Solution {
public:
    int waysToSplit(vector<int>& nums) {
        int n = nums.size();
        vector<long> p(n+1, 0);
        for (int i=1; i<=n; i++) {
            p[i] = p[i-1]+nums[i-1];
        }
        long ans = 0, MOD = 1e9+7;
        for (int i=1; i<n-1; i++) {
            int l = lower_bound(p.begin(), p.end(), 2*p[i])-p.begin();
            l = max(i+1, l);
            int r = upper_bound(p.begin(), p.end(), (p[n]+p[i])/2)-1-p.begin();
            r = min(n-1, r);
            // cout << i << " " <<  l << " " << r << endl;
            ans += max(0, r-l+1);
            ans %= MOD;
        }
        return ans;
    }
};