Prefix sum and edge cases: 1712. Ways to Split Array Into Three Subarrays

1712. Ways to Split Array Into Three Subarrays

Jimmy (xiaoke) Shen

·

Follow

2 min read

·

Jan 3, 2021

--

Listen

Share

O(n²) got TLE

const int MOD = 1e9 + 7;
class Solution {
public:
    int waysToSplit(vector<int>& nums) {
        const int n = nums.size();
        vector<int> pre(n+1, 0);
        for (int i = 1; i <= n; ++i) {
            pre[i] = pre[i-1] + nums[i-1];
        }
        int ret = 0;
        int left, middle, right;
        int j, k;
        for (int i = 0; i < n-2; ++i) {
            left = pre[i+1];
            j = lower_bound(pre.begin(), pre.end(), 2*left) - pre.begin();
            if (j == n + 1) break;
            for (int l = max(i+1, j -1) ; l < n - 1; ++l) {
                middle = pre[l + 1] - pre[i + 1];
                k = lower_bound(pre.begin(), pre.end(), left + 2*middle) - pre.begin();
                if (k == n + 1) break;
                ret++;
                if (ret >= MOD) ret %= MOD;
            }
        }
        return ret;
        
    }
};

O(n log n) got AC

const int MOD = 1e9 + 7;
class Solution {
public:
    int waysToSplit(vector<int>& nums) {
        const int n = nums.size();
        vector<int> pre(n+1, 0);
        for (int i = 1; i <= n; ++i) {
            pre[i] = pre[i-1] + nums[i-1];
        }
        int ret = 0;
        int left, rest;
        // j is the leftmost for the middle
        // k is the leftmost for the right array
        // k - j will be valid solutions.
        int j, k;
        for (int i = 0; i < n-2; ++i) {
            left = pre[i+1];
           
            j = lower_bound(pre.begin(), pre.end(), 2*left) - pre.begin();
            if (j == n + 1) break;
            // the index of leftmost for the middle should be greater than the rightmost of left
            j = max( i +2, j);
            rest = pre[n] - pre[i + 1];
            if (rest % 2 ==0){
               
                k = lower_bound(pre.begin(), pre.end(), pre[n] -rest/2 + 1) - pre.begin();
            } else{
                k = lower_bound(pre.begin(), pre.end(), pre[n]- rest/2) - pre.begin();
            }
            // if k is n + 1, we have an empty right array, this is illegal.
            k = min(n, k);
            if (k <= j) continue;
            ret += k - j;
            if (ret >= MOD) ret %= MOD;
        }
        return ret;
    }
};