旺崽的博客

要么天赋异禀,要么天道酬勤

0%

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

题目链接

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

数组被分成三个 非空 连续子数组,从左至右分别命名为 left , mid , right 。
left 中元素和小于等于 mid 中元素和,mid 中元素和小于等于 right 中元素和。
给你一个 非负 整数数组 nums ,请你返回 好的 分割 nums 方案数目。由于答案可能会很大,请你将结果对 $10^9+7$ 取余后返回。

示例 1:

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

示例 2:

1
2
3
4
5
6
输入: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:

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

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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;
}
};