Storing Values on the Stack When the Answer Is a Distance

A monotonic stack that holds values cannot give a distance, because a distance is the difference between two positions and the value alone does not say where it came from. The tempting repair is to look the popped value up again with Python's index, which returns its first position. That works while every value is distinct, and fails as soon as one repeats: every copy maps to the first.

On 6 4 6 4 9, each position waits for a larger value. The true waits are 4 1 2 1 0. The listing runs the values version with index and the indices version on the same input.

def wait_with_values(values: list[int]) -> list[int]:
    answer = [0] * len(values)
    stack: list[int] = []
    for i, value in enumerate(values):
        while stack and stack[-1] < value:
            waiting = values.index(stack.pop())
            answer[waiting] = i - waiting
        stack.append(value)
    return answer


def wait_with_indices(values: list[int]) -> list[int]:
    answer = [0] * len(values)
    stack: list[int] = []
    for i, value in enumerate(values):
        while stack and values[stack[-1]] < value:
            waiting = stack.pop()
            answer[waiting] = i - waiting
        stack.append(i)
    return answer


values = [6, 4, 6, 4, 9]
print("values on the stack: ", *wait_with_values(values))
print("indices on the stack:", *wait_with_indices(values))

values on the stack:  4 3 0 0 0
indices on the stack: 4 1 2 1 0

The values version prints 4 3 0 0 0. When the 9 pops the second 4, index finds the first 4 at position 1 and overwrites its correct wait of 1 with 3, and the second 4 at position 3 is never set. The same happens to the two 6s: both write to position 0, and position 2 stays 0.

Storing indices fixes both: the index is the position, and the value is one read away as values[stack[-1]]. The index call also costs O(n) per pop, so the values version is O(n²) as well as wrong. When the answer is the value itself, as for the next greater value, values on the stack are enough; when it is a distance or a position, store indices.