I tried following the official solution of LC975 (https://leetcode.com/problems/odd-even-jump/), using a monotonic stack - i.e., I reimplemented their solution (Approach 1) from python in C++.
TLDR Problem statement: Given an array of ints A, "good index" is an index from which the last element of A can be reached, through a series of consecutive "odd" and "even" jumps. Odd jump (1st, 3rd, 5th jump) from index "i" is defined as jump to the smallest element to the right, greater or equal to the current element (A[i]); Even jump (2nd, 4th,...) -> is defined as jump from index "i" to the largest element to the right, smaller or equal to the current element, A[i]. In case of equality, the element with the lowest index is chosen. Return the number of "good" indices.
[Update]: Fixed a bug, passing A by reference instead of value, in the sort lambda -- which fixes the Time Limit Exceeded.
int oddEvenJumps(vector<int>& A) {
int N = static_cast<int>(A.size());
// Compute odd and even jump indices, i.e. where we can jump from here
std::vector<int> odd_next(N, 0);
std::vector<int> even_next(N, 0);
// First, sort indices in A, by value.
for (int ii = 0; ii < N; ++ii) {
odd_next[ii] = ii;
even_next[ii] = ii;
}
// Sort indices in increasing order by value.
std::sort(odd_next.begin(), odd_next.end(), [&A](int a, int b) {
if (A[a] == A[b]) {
return a < b;}
else {
return A[a] < A[b];
}
});
// Sort indices in decreasing order, for even jumps.
//Update: previously -- [A](int a, int b) -- passed A by value -- which was slow.
std::sort(even_next.begin(), even_next.end(), [&A](int a, int b){
if (A[a] == A[b]) {
return a < b;
}
else {
return A[a] > A[b];
}
});
// Convert the sorted indices into next jump index,
// using a monotonic stack.
odd_next = make_vec(odd_next);
even_next = make_vec(even_next);
// Store a flag if from the current index, through an odd jump or
// even jump respectively, we can reach the last element.
vector<int> odd(N, 0);
vector<int> even(N, 0);
odd[N - 1] = 1;
even[N - 1] = 1;
// Go backwards, and mark reachable elements.
for (int ii = N - 2; ii >= 0; --ii) {
if (odd_next[ii] != -1) {
odd[ii] = even[odd_next[ii]];
}
if (even_next[ii] != -1) {
even[ii] = odd[even_next[ii]];
}
}
int res = 0;
for (int jj = 0; jj < N; ++jj){
res += odd[jj];
}
return res;
}
// Use monotonic stack to find the element to jump to, i.e.
// the closest larger or closest smaller element.
inline std::vector<int> make_vec(const std::vector<int>& v){
std::vector<int> result(v.size(), -1);
stack<int> st;
for (int x: v){
while (!st.empty() && x > st.top()) {
result[st.top()] = x;
st.pop();
}
st.push(x);
}
return result;
}
However, this implementation doesn't pass few tests, resulting in time limit exceeded. The complexity of the algorithm should be O(N log N), dominated by the two sort loops.
make_vec
according to my understanding, is O(N), because each element in the vector gets on the stack only once.
How could I improve this code, in particular performance?