Given an integer array nums
of length n
, your task is to count the number of ways to split the array into three non-empty contiguous subarrays such that the sum of the first subarray is less than or equal to the sum of the second, and the sum of the second is less than or equal to the sum of the third.
Constraints:
3 ≤ n ≤ 105
0 ≤ nums[i] ≤ 104
To solve this problem, let's first consider a brute-force approach: Try every possible way to split the array into three contiguous parts, and for each split, check if the subarray sums satisfy the required inequalities. However, with arrays up to 105 elements, this approach is far too slow.
The key insight is to realize that we can use prefix sums to quickly compute subarray sums, and then for each possible position of the first split, efficiently find valid positions for the second split. Rather than checking every possibility, we can use binary search to find the range of valid splits for the second cut, leveraging the fact that the prefix sums are non-decreasing when all numbers are non-negative.
This optimization dramatically reduces the number of checks from O(n2) to O(n log n), making it feasible for large arrays.
Let's break down the optimized solution step by step:
prefix
where prefix[i]
is the sum of nums[0] ... nums[i-1]
(using 0-based indexing).nums[i:j]
in O(1) time as prefix[j] - prefix[i]
.i
(the end of the first subarray), ensure at least one element remains for the second and third subarrays.nums[0:i]
, so i
ranges from 1 to n-2
.i
, we want to find all j
such that:
prefix[i] ≤ prefix[j] - prefix[i]
(first sum ≤ second sum)prefix[j] - prefix[i] ≤ prefix[n] - prefix[j]
(second sum ≤ third sum)j
is the end of the second subarray, and must satisfy i+1 ≤ j ≤ n-1
j
(left bound) where the first inequality holds, and the largest j
(right bound) where the second inequality holds.i
is max(0, right - left + 1)
.i
to get the final answer.This approach ensures that each step is efficient, leveraging prefix sums and binary search to keep the overall time complexity manageable.
Let's use the sample input nums = [1,2,2,2,5,0]
.
prefix = [0,1,3,5,7,12,12]
i
):
i = 1
(first subarray: [1])j
where second subarray is non-empty and both inequalities are satisfied.i=1
:
prefix[1] ≤ prefix[j] - prefix[1]
→ 1 ≤ prefix[j] - 1
→ prefix[j] ≥ 2
prefix[j] - prefix[1] ≤ prefix[6] - prefix[j]
→ prefix[j] - 1 ≤ 12 - prefix[j]
→ 2*prefix[j] ≤ 13
→ prefix[j] ≤ 6
j
must satisfy prefix[j] ≥ 2
and prefix[j] ≤ 6
with 2 ≤ j \leq 5
j
are those where prefix[j]
is 3 or 5 (indices 2, 3), so two ways.i
Values:
i = 2, 3, 4
and sum all valid splits.The final answer is the total number of valid (i, j) pairs found across all possible first cuts.
n
(up to 105).The key to this problem is recognizing that checking every possible split is too slow, and that prefix sums and binary search can be combined to efficiently count valid ways to split the array. By precomputing prefix sums and using binary search for the second cut, we reduce the problem to manageable complexity. This approach is elegant because it leverages the sorted nature of prefix sums (due to non-negative numbers) and classic searching techniques to avoid redundant work.
from bisect import bisect_left, bisect_right
class Solution:
def waysToSplit(self, nums):
n = len(nums)
prefix = [0]
for num in nums:
prefix.append(prefix[-1] + num)
total = prefix[-1]
res = 0
for i in range(1, n - 1):
left = bisect_left(prefix, 2 * prefix[i])
right = bisect_right(prefix, (total + prefix[i]) // 2)
left = max(left, i + 1)
right = min(right, n)
if left <= right:
res += max(0, right - left)
return res
#include <vector>
#include <algorithm>
using namespace std;
class Solution {
public:
int waysToSplit(vector<int>& nums) {
int n = nums.size();
vector<long long> prefix(n + 1, 0);
for (int i = 0; i < n; ++i) {
prefix[i + 1] = prefix[i] + nums[i];
}
int res = 0;
for (int i = 1; i < n - 1; ++i) {
int l = lower_bound(prefix.begin() + i + 1, prefix.end(), 2 * prefix[i]) - prefix.begin();
int r = upper_bound(prefix.begin() + i + 1, prefix.end(), (prefix[n] + prefix[i]) / 2) - prefix.begin();
if (l <= r && l <= n && r <= n) {
res += max(0, r - l);
}
}
return res;
}
};
import java.util.*;
class Solution {
public int waysToSplit(int[] nums) {
int n = nums.length;
long[] prefix = new long[n + 1];
for (int i = 0; i < n; ++i) {
prefix[i + 1] = prefix[i] + nums[i];
}
int res = 0;
for (int i = 1; i < n - 1; ++i) {
int l = Arrays.binarySearch(prefix, i + 1, n + 1, 2 * prefix[i]);
if (l < 0) l = -l - 1;
int r = Arrays.binarySearch(prefix, i + 1, n + 1, (prefix[n] + prefix[i]) / 2);
if (r < 0) r = -r - 1;
else r += 1;
if (l <= r && l <= n && r <= n) {
res += Math.max(0, r - l);
}
}
return res;
}
}
var waysToSplit = function(nums) {
const n = nums.length;
const prefix = [0];
for (let num of nums) {
prefix.push(prefix[prefix.length - 1] + num);
}
let res = 0;
for (let i = 1; i < n - 1; ++i) {
// Binary search for left
let left = i + 1, right = n, l = n, r = n;
let target = 2 * prefix[i];
while (left <= right) {
let mid = Math.floor((left + right) / 2);
if (prefix[mid] >= target) {
l = mid;
right = mid - 1;
} else {
left = mid + 1;
}
}
// Binary search for right
left = i + 1;
right = n;
target = Math.floor((prefix[n] + prefix[i]) / 2);
while (left <= right) {
let mid = Math.floor((left + right) / 2);
if (prefix[mid] > target) {
right = mid - 1;
} else {
r = mid;
left = mid + 1;
}
}
if (l <= r && l <= n && r <= n) {
res += Math.max(0, r - l + 1);
}
}
return res;
};