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
解释：没有好的分割方案。

题意比较简单，就是往一个数组里插两块隔板，暴力的话复杂度 O ( n 2 ) O(n^2) O(n2)，显然不行，考虑优化，很容易想到二分，遍历第一块隔板的位置 i i i，然后找第二块隔板的所有合法位置区间 [ L , R ] [L,R] [L,R] 即可，预处理所有数的前缀和，左端点 L L L 可以通过一个 l o w e r _ b o u n d lower\_bound lower_bound 直接求得，而对右端点 R R R，直接二分求即可，对每一块隔板 i i i，答案加上 R − L + 1 R-L+1 R−L+1，AC代码如下：

class Solution { 
public:
    long long pre[100005] = { 0}, ans = 0, mod = 1e9 + 7, len;
    vector<long long> v;
    
    int bs(int l, int r, int p) { 
        while (l <= r) { 
            int mid = l + r >> 1;
            if (pre[len - 1] - pre[mid] >= pre[mid] - pre[p]) l = mid + 1;
            else r = mid - 1;
        }
        return r;
    }

    int waysToSplit(vector<int> &nums) { 
        len = nums.size();
        for (int i = 0; i < len; i++) { 
            if (i == 0) pre[i] = nums[i];
            else pre[i] = pre[i - 1] + nums[i];
            v.push_back(pre[i]);
        }
        for (int i = 0; i < len - 2; i++) { 
            int pos1 = lower_bound(v.begin() + i + 1, v.end(), pre[i] + pre[i]) - v.begin();
            if (pos1 == v.size()) break;
            int pos2 = bs(pos1, len - 2, i);
            ans = (ans + (long long) (pos2 - pos1 + 1)) % mod;
        }
        return ans;
    }
};