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

Since we need to compute subarray sum, so we should get the prefix sum array ps. Notice that nums[i] is nonnegative, so ps is non-decreasing. Using ps we can either do a linear scan with binary search, or just a linear scan.

Solution 1. O(N*logN), fix the start index of the 3rd subarray and binary search the leftmost and rightmost end index of the 1st subarray.

class Solution {
    public int waysToSplit(int[] nums) {
        int n = nums.length, mod = (int)1e9 + 7, ans = 0;
        int[] ps = new int[n];
        ps[0] = nums[0];
        for(int i = 1; i < n; i++) {
            ps[i] = ps[i - 1] + nums[i];
        }
        //fix the 3rd partition [i, n - 1], its sum * 3 >= total sum
        for(int i = 2; i < n; i++) {
            if((ps[n - 1] - ps[i - 1]) * 3 < ps[n - 1]) {
                break;
            }
            int l = bs1(ps, i - 1, ps[n - 1] - ps[i - 1]);
            int r = bs2(ps, i - 1, ps[n - 1] - ps[i - 1]);
            if(l >= 0 && r >= 0) {
                ans = (ans + r - l + 1) % mod;
            }
        }
        return ans;
    }
    //return the left most index i such that S(nums[0, i]) <= S(nums[i + 1, rightBound]) <= S(nums[rightBound + 1, n - 1])
    private int bs1(int[] ps, int rightBound, int third) {
        int l = 0, r = rightBound - 1;
        while(l < r - 1) {
            int mid = l + (r - l) / 2;
            if(ps[mid] <= ps[rightBound] - ps[mid] && ps[rightBound] - ps[mid] <= third) {
                r = mid;
            }
            //1st is too big, need to keep searching on the left half
            else if(ps[mid] > ps[rightBound] - ps[mid]) {
                r = mid - 1;
            }
            else {
                l = mid + 1;
            }
        }
        if(ps[l] <= ps[rightBound] - ps[l] && ps[rightBound] - ps[l] <= third) {
            return l;
        }
        else if(ps[r] <= ps[rightBound] - ps[r] && ps[rightBound] - ps[r] <= third) {
            return r;
        }
        return -1;
    }
    //return the right most index i such that S(nums[0, i]) <= S(nums[i + 1, rightBound]) <= S(nums[rightBound + 1, n - 1])
    private int bs2(int[] ps, int rightBound, int third) {
        int l = 0, r = rightBound - 1;
        while(l < r - 1) {
            int mid = l + (r - l) / 2;
            if(ps[mid] <= ps[rightBound] - ps[mid] && ps[rightBound] - ps[mid] <= third) {
                l = mid;
            }
            else if(ps[mid] > ps[rightBound] - ps[mid]){
                r = mid - 1;
            }
            else {
                l = mid + 1;
            }
        }
        
        if(ps[r] <= ps[rightBound] - ps[r] && ps[rightBound] - ps[r] <= third) {
            return r;
        }
        else if(ps[l] <= ps[rightBound] - ps[l] && ps[rightBound] - ps[l] <= third) {
            return l;
        }
        return -1;        
    }
}

Solution 2. O(N) fix the end index of the 1st subarray i and find the leftmost and rightmost end index of the 2nd subarray, call them j and k. We actually do not need binary search because as i goes from left to right, the sum of the 1st subarray increases. This means j and k can only either stay the same or increase in each iteration. 

class Solution {
    public int waysToSplit(int[] nums) {
        int n = nums.length, ans = 0;
        for (int i = 1; i < n; ++i)
            nums[i] += nums[i - 1];
        for (int i = 0, j = 0, k = 0; i < n - 2; ++i) {
            while (j <= i || (j < n - 1 && nums[j] < nums[i] * 2))
                j++;
            while (k < j || ( k < n - 1 && nums[k] - nums[i] <= nums[n - 1] - nums[k]))
                k++;
            ans = (ans + k - j) % 1000000007;
        }    
        return ans;
    }    
}