Skip to content

1712. Ways to Split Array Into Three Subarrays

SourceWeekly Contest 222 Q3DifficultyMediumRating2078

Description

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 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

Solutions

Thinking

We need three nonempty parts with \(s_{\textit{left}}\le s_{\textit{mid}}\le s_{\textit{right}}\). Two nested cuts are \(O(n^2)\) and fail for \(n\le 10^5\).

Values are nonnegative, so prefix sums are monotone. After fixing the left cut \(i\), the mid cut lies in a contiguous range and can be found by binary search.

The range is \(s[j]\ge 2s[i]\) and \(s[k]\le (s[-1]+s[i])/2\). Two binary searches per \(i\) count the ways, taken modulo \(10^9+7\).

First, we preprocess the prefix sum array \(s\) of the array \(nums\), where \(s[i]\) represents the sum of the first \(i+1\) elements of the array \(nums\).

Since all elements of the array \(nums\) are non-negative integers, the prefix sum array \(s\) is a monotonically increasing array.

We enumerate the index \(i\) that the left subarray can reach in the range \([0,..n-2)\), and then use the monotonically increasing characteristic of the prefix sum array to find the reasonable range of the mid subarray split by binary search, denoted as \([j, k)\), and accumulate the number of schemes \(k-j\).

In the binary search details, the subarray split must satisfy \(s[j] \geq s[i]\) and \(s[n - 1] - s[k] \geq s[k] - s[i]\). That is, \(s[j] \geq s[i]\) and \(s[k] \leq \frac{s[n - 1] + s[i]}{2}\).

Finally, return the number of schemes modulo \(10^9+7\).

The time complexity is \(O(n \times \log n)\), where \(n\) is the length of the array \(nums\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        mod = 10**9 + 7
        s = list(accumulate(nums))
        ans, n = 0, len(nums)
        for i in range(n - 2):
            j = bisect_left(s, s[i] << 1, i + 1, n - 1)
            k = bisect_right(s, (s[-1] + s[i]) >> 1, j, n - 1)
            ans += k - j
        return ans % mod

 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
31
class Solution {
    private static final int MOD = (int) 1e9 + 7;

    public int waysToSplit(int[] nums) {
        int n = nums.length;
        int[] s = new int[n];
        s[0] = nums[0];
        for (int i = 1; i < n; ++i) {
            s[i] = s[i - 1] + nums[i];
        }
        int ans = 0;
        for (int i = 0; i < n - 2; ++i) {
            int j = search(s, s[i] << 1, i + 1, n - 1);
            int k = search(s, ((s[n - 1] + s[i]) >> 1) + 1, j, n - 1);
            ans = (ans + k - j) % MOD;
        }
        return ans;
    }

    private int search(int[] s, int x, int left, int right) {
        while (left < right) {
            int mid = (left + right) >> 1;
            if (s[mid] >= x) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
public:
    const int mod = 1e9 + 7;

    int waysToSplit(vector<int>& nums) {
        int n = nums.size();
        vector<int> s(n, nums[0]);
        for (int i = 1; i < n; ++i) s[i] = s[i - 1] + nums[i];
        int ans = 0;
        for (int i = 0; i < n - 2; ++i) {
            int j = lower_bound(s.begin() + i + 1, s.begin() + n - 1, s[i] << 1) - s.begin();
            int k = upper_bound(s.begin() + j, s.begin() + n - 1, (s[n - 1] + s[i]) >> 1) - s.begin();
            ans = (ans + k - j) % mod;
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
func waysToSplit(nums []int) (ans int) {
    const mod int = 1e9 + 7
    n := len(nums)
    s := make([]int, n)
    s[0] = nums[0]
    for i := 1; i < n; i++ {
        s[i] = s[i-1] + nums[i]
    }
    for i := 0; i < n-2; i++ {
        j := sort.Search(n-1, func(h int) bool { return h > i && s[h] >= (s[i]<<1) })
        k := sort.Search(n-1, func(h int) bool { return h >= j && s[h] > (s[n-1]+s[i])>>1 })
        ans = (ans + k - j) % mod
    }
    return
}

 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
/**
 * @param {number[]} nums
 * @return {number}
 */
var waysToSplit = function (nums) {
    const mod = 1e9 + 7;
    const n = nums.length;
    const s = new Array(n).fill(nums[0]);
    for (let i = 1; i < n; ++i) {
        s[i] = s[i - 1] + nums[i];
    }
    function search(s, x, left, right) {
        while (left < right) {
            const mid = (left + right) >> 1;
            if (s[mid] >= x) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
    let ans = 0;
    for (let i = 0; i < n - 2; ++i) {
        const j = search(s, s[i] << 1, i + 1, n - 1);
        const k = search(s, ((s[n - 1] + s[i]) >> 1) + 1, j, n - 1);
        ans = (ans + k - j) % mod;
    }
    return ans;
};

Comments