Задача с собеседования в Josh Technology Group

Дан массив целых чисел temperatures, представляющий ежедневные значения температуры. Верните массив answer, где answer[i] - это количество дней, которое нужно подождать после i-ого, чтобы наступил день с более высокой температурой. Если нет будущего дня, для которого это возможно, вместо этого сохраните answer[i] == 0.

Пример 1:
Input: temperatures = [73,74,75,71,69,72,76,73]
Output: [1,1,4,2,1,1,0,0]

Пример 2:
Input: temperatures = [30,40,50,60]
Output: [1,1,1,0]

Пример 3:
Input: temperatures = [30,60,90]
Output: [1,1,0]

Ограничения:
1 <= temperatures.length <= 10⁵
30 <= temperatures[i] <= 100

НАШ ЧАТ АЛГОРИТМИСТОВ

Решение
При наивном решении мы бы итерировались по массиву для каждого дня в поисках более тёплого c асимптотикой O(n²).
Но мы видим паттерн - поиск ближайшего большего/меньшего эл-та, поэтому используем монотонный стек (стек, элементы которого хранятся в строго возрастающем или строго убывающем порядке).
В данном случае стек будет монотонно убывающим. При добавлении нового эл-та алгоритм будет сравнивать его с вершиной стека:
- Пока текущий эл-т больше верхнего эл-та стека (stack[-1][0]): достаём верхний элемент, вычисляем ответ для него (через разницу между индексами текущего эл-та и эл-та из стека) и удаляем эл-т из стека.
Кладём текущий эл-т в стек.


Разберём более подробно:
Создаём:
- стек для хранения пар (температура, индекс) в монотонно убывающем порядке;
- массив answer длиной n, равной длине входящего массива. Заполняем его нулями.

Итерируемся по массиву температур:
Пока стек не пуст и в нём есть дни холоднее текущего:
- достаём значение и индекс более холодного дня, удаляя его из стека;
- вычисляем разницу между индексом текущего дня и индексом более холодного дня - таким образом, узнаем кол-во дней, которые должны пройти между ними. Записываем разницу в массив answer по индексу более холодного дня (answer[stack_i]).
После выхода из цикла while или непопадания в него: добавляем текущий день в стек для последующего сравнения с другими значениями.

В конце возвращаем заполненный массив answer.

Сложность
O(n) - по времени (каждый эл-т добавляется в стек только 1 раз и может быть удалён только 1 раз)
O(n) - по памяти (в худшем случае в стек придётся добавить все элементы входного массива длиной n)

Код
class Solution:
def dailyTemperatures(self, temperatures: List[int]) -> List[int]:
n = len(temperatures)
stack = []
answer = [0] * n

for i, temp in enumerate(temperatures):
while stack and stack[-1][0] < temp:
stack_temp, stack_i = stack.pop()
answer[stack_i] = i - stack_i

stack.append((temp, i))

return answer

@algoses