Gas Station: жадное решение за O(n) и почему старт после провала

Коротко

• Считаем не бензин и расход по отдельности, а разность gas[i] - cost[i]. Это всё, что важно для бака.
• Если суммарный бензин меньше суммарного расхода, ответа нет. Если не меньше, ответ точно есть.
• Если из старта s бак ушёл в минус на станции i, то ни одна станция от s до i не годится в старт. Следующий кандидат — i + 1.
• Ключевой вопрос на собеседовании: «почему можно пропустить все станции между s и i?». Ответ на него и есть доказательство.

Условие

По кругу стоят n заправок. На станции i можно залить gas[i] литров, а дорога от станции i до станции i + 1 стоит cost[i] литров. После последней станции снова идёт нулевая. Бак бесконечный, в начале он пуст. Нужно найти индекс станции, с которой можно проехать полный круг по часовой стрелке и вернуться на неё же. Если такой нет, вернуть -1.

gas  = [1, 2, 3, 4, 5]
cost = [3, 4, 5, 1, 2]
→ 3

gas  = [2, 3, 4]
cost = [3, 4, 3]
→ -1

В первом примере стартуем с третьей станции: заливаем 4, едем за 1, в баке 3. Дальше +5 −2, в баке 6. Потом +1 −3, остаётся 4. Потом +2 −4, остаётся 2. Наконец +3 −5, и в баке ровно 0, но мы уже на месте. Ноль литров на финише — допустимо.

Это задача 134 на LeetCode, уровень Medium. Длина массивов до $10^5$, и условие гарантирует: если ответ существует, он единственный. Последнее нам пригодится.

Переходим к разностям

Бак меняется только на величину gas[i] - cost[i]: заправились, потом потратили на дорогу. Обозначим $d_i = gas_i - cost_i$. Тогда старт s подходит, если все частичные суммы

(индексы по модулю $n$) неотрицательны. Для первого примера разности такие: [-2, -2, -2, 3, 3]. Сразу видно, что первые три станции в старт не годятся: уже на первой дороге бак уходит в минус. Задача свелась к поиску такого циклического сдвига массива, у которого все префиксные суммы не меньше нуля.

Наивное решение: пробуем каждый старт

Самое прямое решение — взять каждую станцию как старт и честно проехать круг. Как только бак ушёл в минус, бросаем этот старт.

def can_complete_circuit_naive(gas: list[int], cost: list[int]) -> int:
    n = len(gas)
    for start in range(n):
        tank = 0
        for k in range(n):
            i = (start + k) % n
            tank += gas[i] - cost[i]
            if tank < 0:
                break
        else:
            return start
    return -1

for ... else в Python срабатывает, если цикл не прервался через break. Это значит, что круг пройден целиком.

Код правильный, но в худшем случае работает за $O(n^2)$. Пример плохого входа: все разности нулевые, кроме последней, равной −1. Тогда каждый старт проезжает почти до конца массива и проваливается только на последней дороге. На моей машине при $n = 10^4$ такой вход обрабатывается около трёх секунд, при $10^5$ было бы около пяти минут. Нужно что-то линейное.

Обратите внимание, сколько работы перебор делает впустую. Стартовали с s, доехали до i и провалились. Затем стартуем с s + 1 и снова проезжаем почти те же станции. Интуиция подсказывает, что второй заход обречён. Это действительно так, и сейчас мы это докажем.

Жадное решение за O(n)

Алгоритм умещается в несколько строк:

def can_complete_circuit(gas: list[int], cost: list[int]) -> int:
    if sum(gas) < sum(cost):
        return -1
    start = 0
    tank = 0
    for i in range(len(gas)):
        tank += gas[i] - cost[i]
        if tank < 0:
            start = i + 1
            tank = 0
    return start


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

Едем от текущего кандидата start и копим бак. Если бак ушёл в минус, объявляем новым кандидатом следующую станцию и обнуляем бак. Проход один, по кругу мы при этом не возвращаемся. Прогон на первом примере:

Суммарная разность $-6 + 6 = 0 \ge 0$, ответ start = 3.

Сумму можно считать в том же цикле, тогда проход по массиву действительно один:

def can_complete_circuit_one_pass(gas: list[int], cost: list[int]) -> int:
    total = 0
    tank = 0
    start = 0
    for i, (g, c) in enumerate(zip(gas, cost)):
        diff = g - c
        total += diff
        tank += diff
        if tank < 0:
            start = i + 1
            tank = 0
    return start if total >= 0 else -1

Выглядит подозрительно просто. Кажется, что мы не проверили две вещи. Почему можно перепрыгнуть сразу на i + 1, не пробуя станции между старым стартом и i? И почему последний start проходит круг, ведь участок от 0 до start - 1 после старта мы вообще не проезжали? Разберём оба вопроса по очереди.

Доказательство

Шаг 1. Если всего бензина не хватает, ответа нет

Это очевидно. За круг мы зальём $\sum gas_i$ литров, а потратим $\sum cost_i$. Если первое меньше, бак где-то обязательно уйдёт в минус, с какой станции ни начинай.

Шаг 2. Почему после провала старт сразу за ним

Лемма. Пусть мы стартовали со станции s и бак впервые ушёл в минус после станции i, то есть $d_s + \ldots + d_i < 0$. Тогда ни одна станция j из отрезка $[s, i]$ не может быть стартом.

Доказательство. Возьмём любую j из $(s, i]$. Раз провал случился впервые на i, на подъезде к j бак был неотрицательным:

Тогда

Значит, если стартовать с j с пустым баком, к станции i мы приедем с отрицательным баком и провалимся не позже. Для самой s это верно по условию. $\blacksquare$

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

Шаг 3. Почему последний кандидат проходит круг

Пусть цикл закончился с кандидатом start, а суммарная разность $T = \sum d_i \ge 0$.

Участок от start до конца массива. После последнего сброса бак ни разу не уходил в минус, иначе сброс был бы позже. Значит, все частичные суммы на этом участке неотрицательны, и в конце массива у нас $S = d_{start} + \ldots + d_{n-1} \ge 0$.

Участок от 0 до start - 1. Его нужно проехать уже с запасом $S$. Возьмём любое $k < start$ и проверим бак на станции k: он равен $S + d_0 + \ldots + d_k$. Отрезок $[0, start - 1]$ целиком разбит на проваленные куски. Сумма каждого куска отрицательна, а по лемме из шага 2 отрицателен и любой хвост куска. Поэтому

а значит,

Для $k = start - 1$ хвост пуст, и бак на финише равен ровно $T \ge 0$. Бак ни разу не ушёл в минус, круг пройден. $\blacksquare$

Попутно мы доказали обратное к шагу 1: если $T \ge 0$, ответ всегда существует. Поэтому в коде достаточно одной проверки суммы и не нужно второй раз проезжать круг от найденного старта.

Другой взгляд: минимум префиксной суммы

Есть ещё одно рассуждение, которое многим кажется нагляднее. Нарисуем график префиксных сумм $P_k = d_0 + \ldots + d_k$ без всякого сброса. Запуск со станции s сдвигает весь график вниз на $P_{s-1}$, то есть на значение в точке перед стартом. Чтобы ни одна точка не ушла ниже нуля, вычитать надо самое маленькое значение графика. Значит, стартовать надо сразу после глобального минимума префиксной суммы.

def can_complete_circuit_min_prefix(gas: list[int], cost: list[int]) -> int:
    prefix = 0
    min_prefix = 0
    min_index = 0
    for i in range(len(gas)):
        prefix += gas[i] - cost[i]
        if prefix < min_prefix:
            min_prefix = prefix
            min_index = i + 1
    return min_index % len(gas) if prefix >= 0 else -1

Для [-2, -2, -2, 3, 3] префиксные суммы равны -2, -4, -6, -3, 0. Минимум −6 достигается на индексе 2, старт — 3. Результат совпадает с жадным алгоритмом, и это не случайность: каждый сброс в жадном решении происходит ровно тогда, когда префиксная сумма опускается ниже предыдущего минимума. Это два описания одного и того же алгоритма. % len(gas) нужен на случай, когда минимум оказался на последней станции: тогда min_index == n, а стартовать надо с нуля.

Если вам ближе такая форма, на собеседовании можно рассказать её. Главное — уметь объяснить, почему она верна. Сама техника префиксных сумм подробно разобрана в отдельной статье.

Сложность

• Наивное решение: $O(n^2)$ времени в худшем случае, $O(1)$ памяти.
• Жадное решение и минимум префикса: один или два прохода, $O(n)$ времени и $O(1)$ дополнительной памяти.

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

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

Проверять только суммарный бак. Условие sum(gas) >= sum(cost) говорит, что ответ существует, но не говорит, какой он. Без прохода со сбросом индекс не найти.

Сдвигать старт на start + 1, а не на i + 1. Получается тот же перебор: после провала мы снова проверяем станции, которые лемма уже вычеркнула. Код останется верным, но на худшем входе вернётся к $O(n^2)$. Ещё и запутаннее, потому что придётся ходить по кругу.

Ходить по кругу в жадном проходе. Иногда пишут цикл до 2 * n с i % n, чтобы «доехать» от найденного старта до него же. Это не ошибка, но и не нужно: шаг 3 доказательства показывает, что проверки суммы хватает.

Сбрасывать при tank <= 0. Ноль в баке — это нормально, ехать дальше можно. При строгом < на примере gas = [1, 1], cost = [1, 1] ответ 0. При <= старт уедет за конец массива, и функция вернёт индекс 2, которого не существует.

Путать направление cost. cost[i] — это дорога от станции i к i + 1. Если прочитать её как дорогу к станции i, разности сдвинутся на одну позицию, и ответ будет неверным.

Не говорить про единственность. Если на интервью спросят «а если стартов несколько?», жадный алгоритм вернёт самый левый из подходящих: все станции левее него вычеркнуты леммой как негодные. В задаче LeetCode такого вопроса не возникает, ответ гарантированно единственный.

Где это встречается и с чем похоже

Gas Station — частый гость в подборках задач на жадные алгоритмы, и разбирают её почти всегда ради доказательства. Код пишется за пару минут, а вот вопрос «почему так можно» отсеивает тех, кто решение просто запомнил. Если вы проговорите лемму о проваленном отрезке, интервьюер обычно остаётся доволен.

Похожие по духу задачи:

• Kadane, максимальная сумма подмассива. Та же механика «накопленная сумма ушла в минус — начинаем заново со следующего элемента», и то же доказательство через отрицательный префикс. На сайте задача есть как «Максимальный подмассив».
• Jump Game. Тоже один проход с одной переменной, и тоже «жадность», которая на деле ничего не выбирает, а только двигает границу. Попробовать можно в задаче «Прыжки по массиву».
• Общие принципы, когда жадный выбор корректен и как это доказывать, собраны в статье про жадные алгоритмы.

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

• Работаем с разностями $d_i = gas_i - cost_i$.
• $\sum d_i < 0$ — ответа нет. $\sum d_i \ge 0$ — ответ есть.
• Провал на станции i при старте с s вычёркивает весь отрезок $[s, i]$: приехать с запасом не хуже, чем стартовать с нулём.
• Один проход со сбросом даёт ответ за $O(n)$ времени и $O(1)$ памяти. Это то же самое, что стартовать сразу после минимума префиксной суммы.