Problem overview
The function next_greater_element receives a vector of integers, nums. For each element it must find the first element to its right that is strictly greater, and report that element's index, not its value. When no element to the right is greater, the entry is -1. The function returns a vector the same length as nums.
In the first example, nums = [1, 3, 2, 4] gives [1, 3, 3, -1]. For the 1, the first greater value to the right is the 3 at index 1. For the 3, it is the 4 at index 3. For the 2, it is also the 4 at index 3. The 4 is last, so it gets -1.
Approach: let elements wait on a stack
The straightforward way is to scan to the right from every element until a greater one appears. For a decreasing vector such as [5, 4, 3, 2, 1] every scan runs to the end, which is O(n^2) in total.
The faster idea turns the problem around. Instead of asking, for each element, where its answer is, go through the vector once and, for each new element, ask which earlier elements it answers. An element whose answer has not been found yet is waiting. The waiting elements are kept on a stack, as indices.
The stack has a useful property. The values at the waiting indices never increase from the bottom to the top. If a waiting value were smaller than the one above it, the later element would have been its answer, and it would no longer be waiting. So when a new element nums[i] arrives, the waiting elements it beats are exactly the ones at the top of the stack, in a row:
- While the top of the stack holds a value smaller than
nums[i], pop it and recordias its answer. - Stop at the first top that is not smaller. By the property above, nothing under it is smaller either.
- Push
i, because it now waits for its own answer.
Whatever is still on the stack at the end never found a greater element to its right, and keeps the -1 it started with.
For [1, 3, 2, 4]: index 0 is pushed. At index 1, the value 3 beats the 1, so the answer for index 0 is 1, and index 1 is pushed. At index 2, the value 2 doesn't beat 3, so index 2 is pushed on top. At index 3, the value 4 beats 2 and then 3, so both get the answer 3. Index 3 is pushed and stays unanswered. The result is [1, 3, 3, -1].
Step by step
- Create the stack as an empty, mutable
Vec<usize>. It holds indices rather than values, because the answer needs indices, and ausizecan indexnumsdirectly. - Create the result as a mutable
Vec<i32>withvec![-1; nums.len()], a vector withnums.len()copies of-1. Every entry starts as "no greater element". - Loop over every index
iin0..nums.len(), left to right. - Pop the elements that
nums[i]beats. Awhileloop runs while the stack isn't empty and the value at the top index is smaller thannums[i]. The top index is read with*stack.last().unwrap():last()gives anOptionholding a reference,unwrap()takes the reference out, and*reads the index. - Take the top off. Inside the loop,
stack.pop().unwrap()removes the top index and stores it intop. It has its answer, so it no longer waits. - Record the answer. Set
result[top]toi as i32. The cast is needed becauseiis ausizeand the vector holdsi32. - Push the current index after the
whileloop, so that it waits for its own answer. - Return
resultas the final expression, with no semicolon.
The solution in Rust
impl Solution {
fn next_greater_element(nums: Vec<i32>) -> Vec<i32> {
let mut stack: Vec<usize> = Vec::new();
let mut result: Vec<i32> = vec![-1; nums.len()];
for i in 0..nums.len() {
while !stack.is_empty() && nums[*stack.last().unwrap()] < nums[i] {
let top = stack.pop().unwrap();
result[top] = i as i32;
}
stack.push(i);
}
result
}
}
The Rust details worth a closer look:
vec![-1; n]. This macro form repeats one valuentimes. The type annotation fixes the elements asi32.- Reading the top.
stack.last()returnsOption<&usize>. On an empty stack it would beNone, andunwrap()would panic. The first part of the condition,!stack.is_empty(), runs first and stops the evaluation with&&, sounwrap()is only reached when the stack has an element. - Why
*is needed. The reference must be dereferenced to get ausizeto put insidenums[...]. The expressionnums[*stack.last().unwrap()]is only a read, and the borrow of the stack ends at once, sostack.pop()is allowed in the loop body. - Two different index types. The stack stores
usizevalues for indexing. The result storesi32values because it may hold-1, which ausizecan't represent. The one cast,i as i32, is where the two meet. - Strict comparison. The test is
<, not<=, so an equal value is not a "greater" element. - Final expression.
resulthas no semicolon, so it is returned. The vector is moved out to the caller.
Complexity
Time: O(n). The loop has an inner while loop, but each index is pushed exactly once and popped at most once. The inner loop therefore runs at most n times in total across the whole call, not n times for each element. Everything else is constant work.
Space: O(n). The stack can hold every index, for example for a decreasing vector. The result vector has n entries.
Edge cases and common mistakes
- Empty vector. The loop doesn't run, and the function returns an empty vector.
- One element.
[5]returns[-1]. - Equal values.
[3, 3, 3]returns[-1, -1, -1], and[2, 2, 5]returns[2, 2, -1]. An equal value is not greater. - Decreasing vector.
[5, 4, 3, 2, 1]returns all-1, and the stack grows to five indices. This is the worst case for memory. - Increasing vector.
[1, 2, 3, 4]returns[1, 2, 3, -1]. Each element is answered by its right neighbour, and the stack never holds more than one index. - Storing values on the stack. The function must report positions. A value would not tell you where the greater element is, and equal values could not be told apart.
- Using
<=. That would treat an equal value as an answer, giving[1, 2, -1]for[3, 3, 3]instead of[-1, -1, -1]. - Forgetting the cast.
result[top] = i;doesn't compile, becauseiis ausizeand the vector holdsi32.
Testing
nums = [1, 3, 2, 4]gives[1, 3, 3, -1].nums = [7, 3, 2, 6, 11, 9, 8, 10, 13]gives[4, 3, 3, 4, 8, 7, 7, 8, -1].nums = [4, 1, 2, 1, 5]gives[4, 2, 4, 4, -1].nums = [2, 2, 5]gives[2, 2, -1].nums = [5]gives[-1].nums = []gives[].
