Problem overview
The function nextGreaterElement receives a slice of integers, nums. For each position it must find the first element to the right that is strictly greater, and return that element's index. When no such element exists, the entry is -1. The result has the same length as nums.
For nums = [1, 3, 2, 4] the answer is [1, 3, 3, -1]. After the 1 comes a 3, at index 1. After the 3, the 2 is smaller and the 4 at index 3 is the first greater value. After the 2, the 4 at index 3 is next. The final 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 indices that are still waiting
The brute-force idea scans to the right from every position until it finds a greater value. That takes O(n²) time on a slice such as a descending one, where no scan ever succeeds.
The faster idea reverses the question. Walk through the slice once, and keep a stack of indices whose next greater element has not been found yet. When a new value nums[i] arrives, it is the answer for every waiting index whose value is smaller. Those indices are answered and removed, and then i joins the stack to wait for its own answer.
The stack stays ordered: the values at the waiting indices never increase from bottom to top. A new value is pushed only after every smaller value above it has been popped, so what remains below is not smaller. This is why the loop can stop at the first top that is not smaller than nums[i]: nothing deeper in the stack can be smaller either. It is called a monotonic stack.
Trace on [1, 3, 2, 4]. Index 0 (value 1) is pushed. Index 1 (value 3) beats the top, so result[0] = 1, and index 1 is pushed. Index 2 (value 2) does not beat 3, so it is pushed on top. Index 3 (value 4) beats 2 and then 3, which sets result[2] = 3 and result[1] = 3, and is pushed. Index 3 never gets an answer, so it keeps its -1.
Step by step
- Create the stack as an empty
[]int. It stores indices, not values. - Create
resultwithmake([]int, len(nums)), then fill it with -1 in a loop overrange result.makegives zeros, and 0 is a valid index, so the -1s must be written by hand. - Loop over the indices of
numswithiin a three-partforloop. - Pop the answered indices. An inner loop with only a condition (Go's
while) runs while the stack is non-empty and the value at the top index is smaller thannums[i]. - Record the answer inside it:
result[stack[len(stack)-1]] = i. Write it before popping, while the top is still there to read. - Pop by reslicing to
stack[:len(stack)-1]. - Push
iwithappend, assigning the result back tostack. Push the index, not the value, since the result needs indices. - Return
result. Indices left on the stack never found a greater element and keep -1.
The solution in Go
func nextGreaterElement(nums []int) []int {
stack := []int{}
result := make([]int, len(nums))
for i := range result {
result[i] = -1
}
for i := 0; i < len(nums); i++ {
for len(stack) > 0 && nums[stack[len(stack)-1]] < nums[i] {
result[stack[len(stack)-1]] = i
stack = stack[:len(stack)-1]
}
stack = append(stack, i)
}
return result
}
The Go details worth a closer look:
makeand zero values.make([]int, n)creates a slice of lengthnfilled with zeros. The loopfor i := range resultyields each index, and assigning -1 to each replaces the zeros.- Go's only loop is
for.for len(stack) > 0 && ... {has just a condition, so it works likewhilein other languages. The three-part form is used for the outer loop. - Nested indexing.
nums[stack[len(stack)-1]]reads the top index from the stack and then the value at that index. The top is the last element, at positionlen(stack)-1. - Short-circuit
&&. The length check comes first, sostack[len(stack)-1]is never evaluated on an empty stack, where it would panic with an index out of range. - Push and pop. Reslicing with
stack[:len(stack)-1]shortens the slice without copying.appendmay move the slice to a larger array, so its result is assigned back tostack.
Complexity
Time: O(n). The inner loop sits inside the outer one, but each index is pushed exactly once and popped at most once. Across the whole run the inner loop body executes at most n times, so the total work is linear, not quadratic.
Space: O(n). The stack can hold every index, for example when the values are non-increasing, and the result holds n entries.
Edge cases and common mistakes
- Empty or single-element slice.
[]gives[]and[5]gives[-1]. The loops simply have nothing to do. - Equal values. The comparison is strict. For
[2, 2, 2]no element is greater than another, so the answer is[-1, -1, -1]. Using<=would pair each 2 with the next one. - Descending input.
[5, 4, 3, 2]gives all -1, and the stack grows to full size. - Ascending input.
[1, 2, 3, 4]gives[1, 2, 3, -1], and the stack never holds more than one index. - Values instead of indices. The answer is the position of the greater element. Storing
nums[i]inresultwould give[3, 4, 4, -1]for the first example. - A skipped pop. Recording the answer without popping would hand the same index a second, wrong answer when a later, even larger value arrives.
- Not skipping a greater value.
[3, 1, 3, 4]gives[3, 2, 3, -1]. The first 3 skips the second 3, which is not greater, and finds the 4.
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].[2, 2, 2]gives[-1, -1, -1].[3, 1, 3, 4]gives[3, 2, 3, -1].[]gives[].