Problem overview
The method nextGreaterElement receives an array 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. When no greater element exists to the right, the entry is -1. The result has 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 last element, 4, has nothing to its right, so its entry is -1. The second example, [7, 3, 2, 6, 11, 9, 8, 10, 13], gives [4, 3, 3, 4, 8, 7, 7, 8, -1].
Approach: a stack of elements still waiting
The direct method scans to the right from every element until it finds a greater one. With a decreasing array such as [5, 4, 3, 2, 1], every scan walks to the end, so the work grows quadratically.
Turn the question around. Instead of asking each element to find its answer, let each new element nums[i] answer every earlier element that is still waiting. Keep a stack of the indices that haven't found a greater element yet. When nums[i] arrives:
- every waiting index whose value is smaller than
nums[i]has just found its answer, namelyi, so record it and remove that index from the stack; - then push
i, which now waits for its own answer.
The stack has a useful shape. Its values never increase from bottom to top: a bigger value would have removed the smaller ones below it. So the top is the smallest value in the stack, and when the top is not smaller than nums[i], nothing below it is either, and the popping can stop. This is called a monotonic stack.
Indices left on the stack when the loop ends never found a greater element. Their entries keep the -1 they started with.
Step by step
- Create the stack.
var stack = [Int]()holds indices that are waiting. - Create the result.
var result = [Int](repeating: -1, count: nums.count)has one entry per element, each starting at -1. It stores indices, not values. - Loop over the positions.
for i in 0..<nums.count. - Answer the waiting elements. Inside, a
whileloop runs as long as the stack is not empty andnums[stack.last!]is smaller thannums[i]. The emptiness check comes first, sostack.last!is safe. - Record the answer. Set
result[stack.last!]toiwhile that index is still on top. - Pop it.
stack.removeLast(). That index has its answer and stops waiting. - Push the current index. After the
whileloop,stack.append(i). Push the index rather than the value, because the index is whatresultneeds. - Return
resultafter the main loop.
The solution in Swift
class Solution {
func nextGreaterElement(_ nums: [Int]) -> [Int] {
var stack = [Int]()
var result = [Int](repeating: -1, count: nums.count)
for i in 0..<nums.count {
while !stack.isEmpty && nums[stack.last!] < nums[i] {
result[stack.last!] = i
stack.removeLast()
}
stack.append(i)
}
return result
}
}
The Swift details worth a closer look:
[Int](repeating:count:). It builds an array of a given length filled with one value. Here it creates the whole result up front, so the code can assign toresult[k]directly. Assigning to an index that doesn't exist yet would crash.stack.last!. The propertylastis an optional:nilfor an empty array. The!force-unwraps it. It is safe only because!stack.isEmptycomes first in the condition, and&&stops at the first false operand, so the unwrap never runs on an empty stack.- Result of
removeLast(). It returns the removed element, but Swift doesn't warn when the call is used as a statement, so the line stands alone. - Value types. The array
numsis passed by value and is never changed.resultandstackare localvararrays. - Strict comparison. The
<means an equal value does not count as greater. In[5, 5, 5]every answer is -1.
Complexity
Time: O(n). The inner loop looks like it could repeat many times for each i, but every index is pushed exactly once and popped at most once. The pops over the whole run therefore number at most n, and the total work is linear.
Space: O(n). The result has n entries, and the stack can hold up to n indices, for example when the array is decreasing.
Edge cases and common mistakes
- Empty array. The loop doesn't run and the method returns an empty array.
- Decreasing array.
[5, 4, 3, 2, 1]gives five entries of -1, and nothing is ever popped. - Increasing array.
[1, 2, 3, 4]gives[1, 2, 3, -1], and the stack holds a single index at a time. - Equal values.
[5, 5, 5]gives[-1, -1, -1], because the next greater element must be strictly greater. - Returning values. The task asks for indices. Storing
nums[i]instead ofiwould give the example[3, 4, 4, -1]. - Popping before recording. After
removeLast(), the top is a different index. Write the answer first, while the right index is on top. - Pushing the value. The stack must hold indices, since the comparison needs
nums[...]and the result needs the position.
Testing
[1, 3, 2, 4]gives[1, 3, 3, -1].[7, 3, 2, 6, 11, 9, 8, 10, 13]gives[4, 3, 3, 4, 8, 7, 7, 8, -1].[]gives[].[5, 5, 5]gives[-1, -1, -1].[1, 2, 3, 4]gives[1, 2, 3, -1].
