1712. Ways to Split Array Into Three Subarrays
A split of an integer array is good if:
- The array is split into three non-empty contiguous subarrays - named
left
,mid
,right
respectively from left to right. - The sum of the elements in
left
is less than or equal to the sum of the elements inmid
, and the sum of the elements inmid
is less than or equal to the sum of the elements inright
.
Given nums
, an array of non-negative integers, return the number of good ways to split nums
. As the number may be too large, return it modulo 109 + 7
.
Example 1:
Input: nums = [1,1,1] Output: 1 Explanation: The only good way to split nums is [1] [1] [1].
Example 2:
Input: nums = [1,2,2,2,5,0] Output: 3 Explanation: There are three good ways of splitting nums: [1] [2] [2,2,5,0] [1] [2,2] [2,5,0] [1,2] [2,2] [5,0]
Example 3:
Input: nums = [3,2,1] Output: 0 Explanation: There is no good way to split nums.
Constraints:
3 <= nums.length <= 105
0 <= nums[i] <= 104
Rust Solution
struct Solution;
const MOD: i64 = 1_000_000_007;
impl Solution {
fn ways_to_split(nums: Vec<i32>) -> i32 {
let n = nums.len();
let mut prefix: Vec<i32> = vec![];
let mut prev = 0;
let mut res = 0;
for i in 0..n {
prev += nums[i];
prefix.push(prev);
}
let total = prev;
let mut mid_start = 0;
let mut mid_end = 0;
for i in 0..n {
let left = prefix[i];
mid_start = mid_start.max(i + 1);
while mid_start < n && prefix[mid_start] - left < left {
mid_start += 1;
}
while mid_end + 1 < n && total - prefix[mid_end] >= prefix[mid_end] - left {
mid_end += 1;
}
if mid_end >= mid_start {
res += (mid_end - mid_start) as i64;
res %= MOD;
}
}
res as i32
}
}
#[test]
fn test() {
let nums = vec![1, 1, 1];
let res = 1;
assert_eq!(Solution::ways_to_split(nums), res);
let nums = vec![1, 2, 2, 2, 5, 0];
let res = 3;
assert_eq!(Solution::ways_to_split(nums), res);
let nums = vec![3, 2, 1];
let res = 0;
assert_eq!(Solution::ways_to_split(nums), res);
let nums = vec![0, 3, 3];
let res = 1;
assert_eq!(Solution::ways_to_split(nums), res);
}
Having problems with this solution? Click here to submit an issue on github.