Daily Temperatures: разбор задачи от перебора за O(n²) до стека индексов за O(n)

Коротко

Это разбор одной конкретной задачи. Общая теория, таблица вариантов и гистограмма лежат в статье про монотонный стек, здесь их не повторяю. Зато здесь есть то, чего там нет: доказательство корректности, строгая оценка сложности и решения без стека.

Условие

Дан массив дневных температур. Для каждого дня нужно сказать, сколько дней ждать более тёплой погоды. Если теплее уже не будет, в ответе 0.

Вход:  [73, 74, 75, 71, 69, 72, 76, 73]
Выход: [ 1,  1,  4,  2,  1,  1,  0,  0]

День с температурой 75 ждёт четыре дня: 71, 69 и 72 не подходят, подходит только 76. На LeetCode это задача 739 уровня medium. Ограничения такие: длина до $10^5$, температуры от 30 до 100. Второе ограничение обычно пропускают, а оно ещё пригодится.

Перед кодом стоит уточнить две вещи. «Теплее» значит строго теплее: день с той же температурой потеплением не считается. А ответ — это расстояние в днях, а не сама температура.

Перебор: честно и медленно

Для каждого дня идём вправо до первого более тёплого.

def daily_temperatures_naive(temperatures: list[int]) -> list[int]:
    n = len(temperatures)
    result = [0] * n

    for i in range(n):
        for j in range(i + 1, n):
            if temperatures[j] > temperatures[i]:
                result[i] = j - i
                break

    return result

Какой вход здесь худший? По привычке отвечают: «строго убывающий массив». В общем случае это так, но при температурах от 30 до 100 строго убывать можно только 71 день. На длинном массиве худший случай другой: все дни одинаковые. Никто не находит потепления, каждый день просматривает весь остаток, и получается $n(n-1)/2$ сравнений. На $10^5$ дней это около $5 \cdot 10^9$.

Вот что показал прогон на ноутбуке:

Цифры зависят от машины, важен порядок. Обратите внимание на последнюю строку: на случайных данных перебор выглядит терпимо, потому что потепление обычно находится быстро. Долго ищут только самые жаркие дни, которым идти до конца. Поэтому такое решение может пройти ваш набор случайных тестов и упасть на тесте из одинаковых значений. Проверяйте его на плоском массиве.

Откуда берётся стек

Посмотрим, что перебор делает лишнего. В примере выше он начинает с дня 71, потом с дня 69, и оба раза заново смотрит на одни и те же следующие дни.

Перевернём взгляд. Не будем для каждого дня искать потепление впереди. Пусть каждый новый день сам смотрит назад и спрашивает: «кого из тех, кто ещё ждёт, я могу обрадовать?» Дни без ответа стоят в очереди ожидания. Пришёл день теплее кого-то из них — этому дню ответ найден, и он уходит.

Осталось понять, какая структура нужна для списка ожидающих. Заметим свойство: температуры ожидающих не растут от старых к новым. Если бы более новый день был теплее старого, старый уже получил бы ответ и не ждал. Значит, самый холодный из ожидающих всегда самый свежий. Новому дню достаточно смотреть на последнего добавленного, а это ровно стек.

Решение

def daily_temperatures(temperatures: list[int]) -> list[int]:
    result = [0] * len(temperatures)  # 0 останется у тех, кто не дождался
    waiting: list[int] = []  # индексы дней без ответа

    for today, temp in enumerate(temperatures):
        while waiting and temp > temperatures[waiting[-1]]:
            day = waiting.pop()
            result[day] = today - day
        waiting.append(today)

    return result


print(daily_temperatures([73, 74, 75, 71, 69, 72, 76, 73]))
# [1, 1, 4, 2, 1, 1, 0, 0]

В стеке лежат индексы, потому что ответ — разность дней. Температуру по индексу всегда можно достать из массива, а наоборот нельзя. После цикла стек специально не разбираем: оставшиеся дни так и не дождались потепления, и у них в result уже стоит 0.

Трассировка на примере

Стек записан как индекс:температура, вершина справа.

В конце в стеке дни 6 и 7, у них остаётся 0. Итог [1, 1, 4, 2, 1, 1, 0, 0] совпадает с ожидаемым.

Посмотрите на день 5. Он снял со стека 69 и 71, но остановился на 75. Дальше можно не проверять: под 75 лежат только дни не холоднее него, значит, и они не холоднее 72.

Почему это правильно

На собеседовании обычно спрашивают только про сложность, но корректность доказывается в три строки, и она хорошо показывает понимание. Инвариант перед каждым днём такой: в стеке лежат ровно те прошедшие дни, для которых потепление ещё не наступило, и их температуры снизу вверх не растут.

1. Снимаем только тех, кого надо. Новый день снимает вершины, пока они холоднее него. Как только вершина не холоднее, все дни под ней тоже не холоднее благодаря порядку в стеке, поэтому проверять их незачем.
2. Ответ минимальный. Пусть день today снял день day. Если бы между ними был день теплее day, он снял бы day раньше, и тот не дожил бы до сегодня. Значит, today и есть первое потепление.
3. Порядок сохраняется. После снятий на вершине либо ничего нет, либо день не холоднее сегодняшнего. Кладём сегодняшний, и порядок не нарушен.

Отсюда видно, зачем нужно строгое >. Равные дни не снимают друг друга и лежат в стеке рядом, поэтому порядок нестрогий: не растёт, а не убывает строго.

Почему это O(n), хотя внутри while

Отдельный шаг цикла бывает дорогим. На входе [50, 50, 50, 50, 51] четыре дня ничего не снимают, а последний снимает сразу всех. В общем случае один шаг может стоить $O(n)$, и умножение «$n$ шагов на $O(n)$» даёт оценку $O(n^2)$. Оценка верная, но грубая: дорогие шаги не могут идти подряд. Чтобы снять элемент, его сначала надо положить, и снять его можно только один раз.

Подсчёт по элементам

Считаем не по шагам, а по индексам. Каждый индекс один раз входит в стек и не больше одного раза выходит. Значит, за весь проход:

• append выполняется ровно $n$ раз;
• pop выполняется не больше $n$ раз;
• условие while проверяется не больше $2n$ раз: каждая истинная проверка заканчивается pop, а ложная бывает одна на день.

На примере из трассировки: 8 добавлений, 6 снятий, 2 индекса остались в стеке.

Оценку можно уточнить: сравнений температур не больше $2n-3$ при $n \ge 2$. Снятий не больше $n-1$, потому что последний день всегда остаётся в стеке. Неудачных сравнений тоже не больше $n-1$, потому что в день 0 сравнивать не с кем. Оба максимума сразу невозможны: если снято $n-1$ дней, то в последний день стек опустел, и неудачного сравнения в этот день не было. Граница достигается на входе вида «много одинаковых дней и один теплее в конце».

То же через потенциал

Формальная версия того же рассуждения. Возьмём потенциал $\Phi_i$, равный размеру стека после дня $i$. Пусть в день $i$ снято $k_i$ элементов, тогда реальная стоимость шага равна $c_i = k_i + 1$: столько снятий и одно добавление. Размер стека меняется на $\Phi_i - \Phi_{i-1} = 1 - k_i$. Амортизированная стоимость шага оказывается постоянной:

Сумма по всем дням телескопируется, а изначально стек пуст, поэтому $\Phi_0 = 0$:

Получилось точное равенство: число операций со стеком равно $2n$ минус размер того, что осталось в конце. В примере это $16 - 2 = 14$, ровно те самые 8 добавлений и 6 снятий.

Смысл потенциала простой. Каждое добавление платит за себя и заранее оплачивает своё будущее снятие. Дорогой шаг просто тратит то, что накопили предыдущие.

Две важные оговорки

Амортизация не про среднее. Это гарантия для всей последовательности операций на любом входе, случайность здесь ни при чём. Разницу иногда проверяют отдельным вопросом, подробнее в статье про вопросы о сложности на собеседовании.

list.append в Python тоже амортизированный. Отдельное добавление иногда копирует весь массив, но $n$ добавлений всё равно стоят $O(n)$, это разобрано в статье про динамический массив. Итоговая оценка не меняется.

Память: $O(n)$ в худшем случае. При ограничениях задачи стек забивается целиком снова на одинаковых температурах, ведь никто никого не снимает.

Типичные ошибки

>= вместо >. На [70, 70, 71] правильный ответ [2, 1, 0], а с нестрогим сравнением выходит [1, 1, 0]: первый день считает второй потеплением. На массивах без повторов ошибка не видна, поэтому тест с равными днями обязателен.

if вместо while. Один новый день может закрыть сразу несколько ожидающих. На [73, 72, 71, 75] правильно [3, 2, 1, 0], а с if получится [0, 0, 1, 0].

Добавление только в ветке else. Сегодняшний день нужно класть в стек всегда, даже если он кого-то снял: ему тоже нужен ответ. На [70, 75, 80] такой код вернёт [1, 0, 0] вместо [1, 1, 0].

Шаблон «следующего большего» без правок. Он инициализирует ответ числом -1 и записывает значение. Здесь по умолчанию 0, а записывать надо today - day.

Значения в стеке вместо индексов. Без индекса расстояние не посчитать. Придётся хранить пары, а это лишний код.

Без стека: два решения для вопроса «а можно иначе?»

Оба идут справа налево и не нужны как первый ответ. Зато они хорошо подходят для дополнительных вопросов интервьюера. Заодно проверяют, понимаете ли вы, откуда берётся линейность.

Используем диапазон температур

Температур всего 71 вариант. Идём с конца и для каждой температуры помним ближайший день, в который она была.

def daily_temperatures_buckets(temperatures: list[int]) -> list[int]:
    n = len(temperatures)
    result = [0] * n
    nearest = [n] * 101  # nearest[t] — ближайший день правее с температурой t

    for i in range(n - 1, -1, -1):
        temp = temperatures[i]
        warmer = min(nearest[temp + 1 :], default=n)
        if warmer < n:
            result[i] = warmer - i
        nearest[temp] = i

    return result

Время $O(n \cdot W)$, где $W$ — размер диапазона, дополнительная память $O(W)$. При $W = 71$ это честная константа, и решение проходит. Но оно перестаёт работать, как только температуры становятся произвольными, например вещественными или с большим разбросом. Это плата за использование ограничения.

Прыжки по готовым ответам

Идём справа налево. Когда считаем день i, ответы для всех дней правее уже есть. Смотрим на следующий день: если он теплее, ответ готов. Если нет, то все дни до его собственного потепления тоже не теплее него, поэтому их можно перепрыгнуть.

def daily_temperatures_jumps(temperatures: list[int]) -> list[int]:
    n = len(temperatures)
    result = [0] * n
    hottest = float("-inf")  # максимум среди дней правее

    for i in range(n - 1, -1, -1):
        if temperatures[i] >= hottest:
            hottest = temperatures[i]  # теплее не будет, остаётся 0
            continue

        j = i + 1
        while temperatures[j] <= temperatures[i]:
            j += result[j]
        result[i] = j - i

    return result

На примере: день 2 с температурой 75 смотрит на день 3 (71), прыгает на result[3] = 2 дня вперёд, попадает на день 5 (72), прыгает ещё на 1 и оказывается на дне 6 (76). Ответ $6-2=4$.

Проверка hottest защищает от зацикливания. Без неё день без потепления упёрся бы в result[j] = 0 и прыгал бы на месте.

Вложенный цикл здесь снова линейный, и причина та же. Цепочка прыжков от i + 1 проходит ровно по тем дням, которые лежали бы в стеке при обходе справа налево: стек никуда не делся, просто хранится внутри result. День, через который перепрыгнули, больше никто не посетит, потому что следующая цепочка начнётся с i и сразу пролетит мимо него. На один день приходится не больше одного такого «перепрыгивания» и не больше одной успешной проверки, то есть всего не больше $2n$.

Сравнение подходов:

У стека есть преимущество, которого нет у остальных: он работает слева направо и не требует знать будущее. Если температуры приходят потоком, ответ для дня выдаётся в момент, когда наступило потепление. Решениям справа налево нужен весь массив сразу.

Родственные задачи

Поменяв пару строк, вы решите ещё несколько задач. Вот что меняется в каждой:

Next Greater Element I

Даны два массива: nums1 — подмножество nums2. Для каждого элемента nums1 нужно найти следующий больший в nums2, а если его нет, вернуть -1. Расстояние не нужно, элементы nums2 по условию уникальны, поэтому это редкий случай, когда в стеке можно хранить значения.

def next_greater_element(nums1: list[int], nums2: list[int]) -> list[int]:
    next_greater: dict[int, int] = {}
    stack: list[int] = []

    for num in nums2:
        while stack and num > stack[-1]:
            next_greater[stack.pop()] = num
        stack.append(num)

    return [next_greater.get(num, -1) for num in nums1]


print(next_greater_element([4, 1, 2], [1, 3, 4, 2]))  # [-1, 3, -1]

Next Greater Element II

Массив кольцевой: после последнего элемента снова идёт первый. Стандартный приём — пройти два круга, но класть в стек только на первом. На втором круге элементы лишь закрывают тех, кто ещё ждёт.

def next_greater_elements(nums: list[int]) -> list[int]:
    n = len(nums)
    result = [-1] * n
    stack: list[int] = []

    for i in range(2 * n):
        num = nums[i % n]
        while stack and num > nums[stack[-1]]:
            result[stack.pop()] = num
        if i < n:
            stack.append(i)

    return result


print(next_greater_elements([1, 2, 1]))  # [2, -1, 2]

Амортизация здесь работает так же: проходов $2n$, а добавлений и снятий по-прежнему не больше $n$ каждого.

Online Stock Span

Цены приходят по одной. На каждую нужно вернуть размах: сколько дней подряд, считая сегодняшний и идя назад, цена была не выше сегодняшней. Это зеркальная задача: ищем предыдущий строго больший, причём без доступа ко всему массиву.

class StockSpanner:
    def __init__(self) -> None:
        self.stack: list[tuple[int, int]] = []  # пары (цена, размах)

    def next(self, price: int) -> int:
        span = 1
        while self.stack and self.stack[-1][0] <= price:
            span += self.stack.pop()[1]
        self.stack.append((price, span))
        return span


spanner = StockSpanner()
print([spanner.next(p) for p in [100, 80, 60, 70, 60, 75, 85]])
# [1, 1, 1, 2, 1, 4, 6]

Здесь снятый день схлопывается в размах того, кто его снял, поэтому индексы не нужны. Сравнение нестрогое, потому что равная цена входит в размах. И это лучший пример, зачем нужна амортизация: один вызов next может стоить $O(n)$, но $n$ вызовов вместе стоят $O(n)$, то есть $O(1)$ на вызов амортизированно.

Как рассказывать на собеседовании

Точной статистики по компаниям я привести не могу: публичные списки собраны из отзывов и быстро устаревают. Надёжно одно: задача известная и входит в раздел про стек популярной подборки NeetCode. По сути она проверяет, знаком ли кандидат с монотонным стеком.

Поэтому код сам по себе мало что покажет, важнее объяснение. Рассуждать вслух можно примерно так:

Перебор даёт $O(n^2)$, худший случай — одинаковые температуры. Заведу стек индексов дней, которые ещё ждут потепления. Температуры в нём не растут, поэтому каждому новому дню достаточно смотреть на вершину. Сложность линейная: каждый индекс входит и выходит не больше одного раза, всего не больше $2n$ операций. Память $O(n)$.

После этого стоит прогнать три теста: один день ([50] → [0]), одинаковые дни ([50, 50] → [0, 0]), равные дни перед потеплением ([70, 70, 71] → [2, 1, 0]).

Вероятные дополнительные вопросы:

• «Можно ли без дополнительной памяти?» Прыжки по result.
• «Температуры приходят потоком?» Стек подходит как есть, ответ выдаётся при снятии. Память будет расти на длинных периодах без потепления.
• «Нужен не ближайший тёплый день, а самый тёплый впереди?» Стек не нужен: хватит прохода справа налево с текущим максимумом.
• «А если нужен предыдущий более тёплый день?» Тот же проход, но снимаем при >=, а ответ читаем с вершины сразу после снятий.

Что запомнить

• Стек хранит индексы дней без ответа, снизу вверх температуры не растут.
• Снятие дня и есть момент, когда для него найден ответ.
• Сравнение строгое, while, а не if, и текущий день кладём в стек всегда.
• Операций со стеком ровно $2n$ минус то, что осталось в конце, поэтому $O(n)$.
• Проверяйте решение на массиве из одинаковых значений: он ловит и квадратичное время, и ошибку в сравнении.

Теперь решите задачу сами на JavaScript, Python или Java, не подглядывая в код выше. Затем попробуйте Next Greater Element II: если получится без подсказок, технику вы усвоили.