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

--

--

--

Data Scientist/MLE/SWE @takemobi

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

Recommended from Medium

Importing multiple files from multiple directories in python.

My Experience with Passing the Google Cloud Professional Machine Learning Engineer Certificate Exam

Agile Scrum (OneID)

Congratulations, you graduated your online coding course! What now…?

Using the Needleman-Wunsch algorithm to draw evolutionary trees

PYTHON PROGRAMMING: TOP 12 APPLICATIONS OF PYTHON IN BUSINESS

Symfony messenger and AWS SNS/SQS

How to Create an AWS account?

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store
Jimmy Shen

Jimmy Shen

Data Scientist/MLE/SWE @takemobi

More from Medium

Tips for Job Interview Code Challenges

Know how to get SDE Internship during college irrespective of your branch

Improve your deep learning performance using BlockDrop

ACID(SQL) vs BASE(No-SQL) properties

ACID vs BASE properties (SQL vs NoSQL)