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

## 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;    }};`

--

--

--

## More from Jimmy Shen

Data Scientist/MLE/SWE @takemobi

Love podcasts or audiobooks? Learn on the go with our new app.

## Jimmy Shen

Data Scientist/MLE/SWE @takemobi