leetcode - 1712. Ways to Split Array Into Three Subarrays

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

Solution

Binary Search

Get a pre sum array first, then let’s say the nums is splitted into 3 parts, nums[0:i], nums[i:j], nums[j:]. To satisfy the requirements, we need:
pre_sum[i] <= pre_sum[j] - pre_sum[i] <= pre_sum[-1] - pre_sum[j]
In other words, for every i, we have this range for j:

• pre_sum[j] >= 2 * pre_sum[i]
• pre_sum[j] * 2 <= pre_sum[-1] - pre_sum[i]

Since pre_sum is a sorted array, we could use binary search to find j.

Time complexity: $o(n\log n)$
Space complexity: $o(n)$

Code

Binary Search

class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        def find_smallest(nums: list, target: int) -> int:
            """
            Find smallest index that nums[i] >= target
            """
            left, right = 0, len(nums) - 1
            while left < right:
                mid = (left + right) >> 1
                if nums[mid] < target:
                    left = mid + 1
                else:
                    right = mid
            return (left + right) >> 1
        def find_largest(nums: list, target: int) -> int:
            """
            Find largest index that nums[i] <= target
            """
            left, right = 0, len(nums) - 1
            while left < right:
                mid = (left + right + 1) >> 1
                if nums[mid] > target:
                    right = mid - 1
                else:
                    left = mid
            return (left + right) >> 1
        pre_sum = []
        cur_sum = 0
        for each_num in nums:
            cur_sum += each_num
            pre_sum.append(cur_sum)
        res = 0
        mod_num = 1000000007
        for i in range(len(pre_sum) - 2):
            j_min = find_smallest(pre_sum, pre_sum[i] * 2)
            j_max = find_largest(pre_sum, (pre_sum[-1] + pre_sum[i]) / 2)
            if pre_sum[j_min] >= 2 * pre_sum[i] and pre_sum[j_max] * 2 <= pre_sum[-1] + pre_sum[i]:
                res += max(0, min(len(pre_sum) - 2, j_max) - max(i + 1, j_min) + 1)
                res %= mod_num
        return res % mod_num