Library · Arrays · monotonic stack
Monotonic Stack
# beats(x, y) says x wins against y: operator.gt for "greater", operator.lt for "smaller"
# next_index(a, gt) next greater next_index(a, lt) next smaller
# prev_index(a, gt) previous greater prev_index(a, lt) previous smaller
# use ge / le instead to count ties as wins
# nxt[i] = first j > i with beats(a[j], a[i]), or n
def next_index(a: list, beats) -> list:
n = len(a)
nxt = [n] * n
st = [] # indices still waiting for their answer
for i in range(n):
while st and beats(a[i], a[st[-1]]):
nxt[st.pop()] = i
st.append(i)
return nxt
# prv[i] = last j < i with beats(a[j], a[i]), or -1
def prev_index(a: list, beats) -> list:
n = len(a)
prv = [-1] * n
st = [] # indices that can still be someone's answer
for i in range(n):
while st and not beats(a[st[-1]], a[i]):
st.pop()
if st: prv[i] = st[-1]
st.append(i)
return prv