Problem overview
The method nextGreaterElement receives an Array[Int] called nums. For every position, it looks to the right for the first element that is strictly greater, and records that element's index. If there is none, it records -1. The result is an Array[Int] of the same length as nums.
For nums = [1, 3, 2, 4] the result is [1, 3, 3, -1]: for the 1 the first greater element is the 3 at index 1, for the 3 it is the 4 at index 3, for the 2 it is also the 4, and the last 4 has nothing to its right.
Approach: keep the unanswered indices on a stack
The brute force scans to the right from every position, which is O(n²) for a descending array. The stack avoids repeating the scans.
Walk through the array from left to right and keep a stack of the indices that are still waiting for an answer. When you reach index i, look at the waiting index on top of the stack. If nums(i) is greater than the value there, then i is its next greater element, because it is the first element to the right that beats it. Record the answer and pop it, then check the new top, and keep going until the top isn't beaten. Finally push i, which now waits for its own answer.
The values at the waiting indices never increase from the bottom of the stack to the top. A larger value arriving would have popped every smaller one above it, and an equal value is not popped, since equal is not greater. That is why the loop can stop at the first top that isn't beaten: nothing below it is beaten either. Indices that stay on the stack until the end never found a greater element, and they keep the -1 they started with.
Here is [1, 3, 2, 4], with the stack shown top first:
i=0 (1): stack is empty, push 0 stack: 0
i=1 (3): 3 > 1, result(0) = 1, pop; push 1 stack: 1
i=2 (2): 2 < 3, nothing to pop; push 2 stack: 2, 1
i=3 (4): 4 > 2, result(2) = 3, pop
4 > 3, result(1) = 3, pop; push 3 stack: 3
The result is [1, 3, 3, -1], with index 3 still holding its starting -1.
Step by step
- Create the stack. Declare
stackas avalholding an emptyscala.collection.mutable.Stack[Int](). It holds the indices of elements that are still waiting. - Create the result.
Array.fill(nums.length)(-1)makes an array of the same length filled with-1, the answer for an element that never finds a greater one.filltakes the length and the value in two separate parameter lists. - Loop over the indices:
for (i <- nums.indices). - Pop the beaten indices. A
whileloop runs as long asstack.nonEmpty && nums(stack.top) < nums(i). The first condition protectsstack.topfrom an empty stack. - Record the answer. Inside the loop, set
result(stack.top) = i. Do this before popping, while the index is still on top. - Pop with
stack.pop(). That index has its answer and stops waiting. - Push the current index with
stack.push(i), after thewhileloop. Push the index, not the value, becauseresultstores indices. - Return
resultas the last expression of the method.
The solution in Scala
class Solution {
def nextGreaterElement(nums: Array[Int]): Array[Int] = {
val stack = scala.collection.mutable.Stack[Int]()
val result = Array.fill(nums.length)(-1)
for (i <- nums.indices) {
while (stack.nonEmpty && nums(stack.top) < nums(i)) {
result(stack.top) = i
stack.pop()
}
stack.push(i)
}
result
}
}
The Scala details worth a closer look:
- Arrays use parentheses.
nums(i)reads an element andresult(stack.top) = iwrites one. Square brackets in Scala are for type parameters, as inStack[Int]. topandnonEmptytake no parentheses, whilepop()andpush(i)are called with them.Array.fill(n)(-1)creates anArray[Int]because the value-1is anInt.- Short-circuit
&&. When the stack is empty, the right side is never evaluated, sostack.topis never called on an empty stack. Calling it would throwNoSuchElementException. - No
returnneeded. The last lineresultis the method's value.
Complexity
Time: O(n). The inner while loop sits inside the for loop, but every index is pushed exactly once and popped at most once. The total number of pops across the whole run is at most n, so the total work is O(n).
Space: O(n). The stack can hold every index, for example in a descending array, and the result array has n entries.
Edge cases and common mistakes
- Empty array. The loop doesn't run and the method returns an empty array.
- One element.
[5]returns[-1]. - Equal values. The comparison is strict, so
[2, 2]returns[-1, -1]. Using<=would wrongly answer[1, -1]. - Descending array.
[5, 4, 3, 2, 1]returns all-1and ends with every index on the stack. - Ascending array.
[1, 2, 3, 4, 5]returns[1, 2, 3, 4, -1]: each element pops the one before it. - A mix.
[2, 1, 2, 4, 3]returns[3, 2, 3, -1, -1]. - Storing values instead of indices. The answer is a position. Writing
nums(i)intoresultwould return the greater values, not where they are. - Popping before recording. After
stack.pop(),stack.toprefers to the next index down, so the answer would be written to the wrong place. - Using
ifinstead ofwhile. In the trace, the 4 resolves two indices at once. Anifwould resolve only one.
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 = [2, 2]gives[-1, -1].nums = [2, 1, 2, 4, 3]gives[3, 2, 3, -1, -1].nums = []gives[].
