Ways to Split Array Into Three Subarrays

Problem

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 in mid, and the sum of the elements in mid is less than or equal to the sum of the elements in right.

Given nums, an array of non-negative integers, return the number ofgood ways to split nums. As the number may be too large, return it modulo 10^9 + 7.

Examples

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 <= 10^5
• 0 <= nums[i] <= 10^4

Solution

Method 1 – Prefix Sum + Binary Search

Intuition

We want to split the array into three contiguous parts such that sum(left) ≤ sum(mid) ≤ sum(right). By using prefix sums, we can efficiently compute subarray sums and use binary search to find valid split points.

Approach

1. Compute prefix sums of nums.
2. For each possible end of left part (i), use binary search to find:
   • The smallest j (start of mid) such that sum(left) ≤ sum(mid).
   • The largest j such that sum(mid) ≤ sum(right).
3. For each i, add the number of valid j to the answer.
4. Return the answer modulo 1e9+7.

Code

C++

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

Go

import "sort"
func waysToSplit(nums []int) int {
    const mod = 1_000_000_007
    n := len(nums)
    pre := make([]int, n+1)
    for i := 0; i < n; i++ {
        pre[i+1] = pre[i] + nums[i]
    }
    ans := 0
    for i := 1; i <= n-2; i++ {
        l := sort.Search(n-i, func(j int) bool { return pre[i+j] >= 2*pre[i] }) + i
        r := sort.Search(n-i, func(j int) bool { return pre[i+j] > (pre[n]+pre[i])/2 }) + i
        if l < r {
            ans = (ans + r - l) % mod
        }
    }
    return ans
}

Java

import java.util.*;
class Solution {
    public int waysToSplit(int[] nums) {
        int n = nums.length, mod = 1_000_000_007;
        int[] pre = new int[n+1];
        for (int i = 0; i < n; i++) pre[i+1] = pre[i] + nums[i];
        int ans = 0;
        for (int i = 1; i <= n-2; i++) {
            int l = lowerBound(pre, i+1, n-1, 2*pre[i]);
            int r = lowerBound(pre, i+1, n-1, (pre[n]+pre[i])/2+1);
            if (l < r) ans = (ans + r - l) % mod;
        }
        return ans;
    }
    private int lowerBound(int[] a, int l, int r, int x) {
        while (l < r) {
            int m = (l + r) / 2;
            if (a[m] < x) l = m + 1;
            else r = m;
        }
        return l;
    }
}

Kotlin

class Solution {
    fun waysToSplit(nums: IntArray): Int {
        val n = nums.size
        val mod = 1_000_000_007
        val pre = IntArray(n+1)
        for (i in 0 until n) pre[i+1] = pre[i] + nums[i]
        var ans = 0
        for (i in 1..n-2) {
            val l = lowerBound(pre, i+1, n, 2*pre[i])
            val r = lowerBound(pre, i+1, n, (pre[n]+pre[i])/2+1)
            if (l < r) ans = (ans + r - l) % mod
        }
        return ans
    }
    private fun lowerBound(a: IntArray, l: Int, r: Int, x: Int): Int {
        var left = l
        var right = r
        while (left < right) {
            val m = (left + right) / 2
            if (a[m] < x) left = m + 1 else right = m
        }
        return left
    }
}

Python

from typing import List
class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        from bisect import bisect_left, bisect_right
        n = len(nums)
        pre = [0]
        for x in nums:
            pre.append(pre[-1] + x)
        ans = 0
        mod = 10**9 + 7
        for i in range(1, n-1):
            l = bisect_left(pre, 2*pre[i], i+1, n)
            r = bisect_right(pre, (pre[n]+pre[i])//2, i+1, n)
            if l < r:
                ans = (ans + r - l) % mod
        return ans

Rust

impl Solution {
    pub fn ways_to_split(nums: Vec<i32>) -> i32 {
        let n = nums.len();
        let mut pre = vec![0; n+1];
        for i in 0..n { pre[i+1] = pre[i] + nums[i]; }
        let mut ans = 0;
        let m = 1_000_000_007;
        for i in 1..n-1 {
            let l = (i+1..n).find(|&j| pre[j] >= 2*pre[i]).unwrap_or(n);
            let r = (i+1..n).rfind(|&j| pre[j] <= (pre[n]+pre[i])/2).map_or(i+1, |x| x+1);
            if l < r { ans = (ans + r - l) % m; }
        }
        ans
    }
}

TypeScript

class Solution {
    waysToSplit(nums: number[]): number {
        const n = nums.length, mod = 1_000_000_007;
        const pre = [0];
        for (const x of nums) pre.push(pre[pre.length-1] + x);
        let ans = 0;
        for (let i = 1; i <= n-2; i++) {
            let l = this.lowerBound(pre, 2*pre[i], i+1, n);
            let r = this.upperBound(pre, (pre[n]+pre[i])>>1, i+1, n);
            if (l < r) ans = (ans + r - l) % mod;
        }
        return ans;
    }
    lowerBound(a: number[], x: number, l: number, r: number): number {
        while (l < r) {
            const m = (l + r) >> 1;
            if (a[m] < x) l = m + 1;
            else r = m;
        }
        return l;
    }
    upperBound(a: number[], x: number, l: number, r: number): number {
        while (l < r) {
            const m = (l + r) >> 1;
            if (a[m] <= x) l = m + 1;
            else r = m;
        }
        return l;
    }
}

Complexity

• ⏰ Time complexity: O(n log n) — For each split point, we use binary search, and there are O(n) split points.
• 🧺 Space complexity: O(n) — For the prefix sum array.