Task Scheduler (LeetCode 621): формула, куча и очередь кулдауна

Коротко

Задача 621 на LeetCode — классика уровня Medium, которая часто всплывает на собеседованиях как проверка жадного мышления. Её можно решить двумя способами: красивой формулой в три строки или честной симуляцией. Хорошо бы знать оба — и уметь объяснить, почему формула вообще работает.

Условие

Дан массив задач tasks, каждая обозначена заглавной буквой, и число n. Процессор за один такт выполняет одну задачу или простаивает (idle). Между двумя одинаковыми задачами должно пройти не меньше n тактов. Порядок выполнения любой. Нужно найти минимальное общее число тактов.

tasks = ["A","A","A","B","B","B"], n = 2
Ответ: 8
Расписание: A B idle A B idle A B

tasks = ["A","C","A","B","D","B"], n = 1
Ответ: 6
Расписание: A B A B C D — простоев нет

tasks = ["A","A","A","B","B","B"], n = 3
Ответ: 10
Расписание: A B idle idle A B idle idle A B

Интуиция: самая частая задача задаёт каркас

Представьте, что задача A встречается чаще всех — f раз. Между каждой парой соседних A должно уместиться минимум n тактов. Значит, расписание не может быть короче, чем:

A _ _ A _ _ A        (f = 3, n = 2)

Здесь f - 1 = 2 блока длины n + 1 = 3 и ещё одна A в конце. Итого $(f - 1)(n + 1) + 1$ тактов. Пустые места — это слоты, куда можно положить другие задачи. Если не хватит, будут простои.

Теперь добавим B, которая встречается столько же раз, сколько A. Её последняя копия тоже не влезет в блоки: в каждом блоке может стоять только одна B, а блоков f - 1. Значит, хвост удлиняется:

A B _ A B _ A B

Отсюда «+ k» в формуле: $k$ — число задач с максимальной частотой. Все они дают по одной копии в хвост.

Когда простоев нет совсем

Если разных задач много, слоты в блоках заполнятся полностью, и задачи начнут «выпирать» за каркас. Например, AAABBBCCDDEE при n = 2: каркас даёт $(3 - 1) \cdot 3 + 2 = 8$, а задач 12. Но 12 тактов меньше быть не может никогда — каждую задачу нужно выполнить хотя бы раз. И 12 тактов всегда достижимо: блоки просто расширяются, а вставленные в них дополнительные задачи только увеличивают расстояние между одинаковыми.

Поэтому ответ — максимум из двух нижних оценок:

Решение формулой

from collections import Counter


def least_interval(tasks: list[str], n: int) -> int:
    counts = Counter(tasks).values()
    max_freq = max(counts)
    # сколько задач делят максимальную частоту
    max_count = sum(1 for c in counts if c == max_freq)
    return max(len(tasks), (max_freq - 1) * (n + 1) + max_count)


least_interval(list("AAABBB"), 2)  # 8
least_interval(list("ACABDB"), 1)  # 6
least_interval(list("AAABBB"), 3)  # 10

Время — $O(m)$, где $m$ — число задач: один проход на подсчёт. Память — $O(1)$, потому что различных букв не больше 26. Если задачи — произвольные строки, память $O(k)$ по числу различных задач.

Почему формула — это не просто нижняя оценка

На собеседовании мало сказать «ответ не меньше X». Нужно показать, что X достижим. Аргумент такой:

1. Выписываем f - 1 блоков, в начале каждого — по одной копии каждой задачи с максимальной частотой, плюс хвост из этих же задач.
2. Остальные задачи раскладываем по блокам по кругу: первую копию в первый блок, вторую во второй и так далее. Задача встречается не больше f - 1 раз (иначе она тоже была бы максимальной), поэтому две её копии никогда не попадут в один блок.
3. Если пустых слотов хватило — ответ равен каркасу, простои заполнены не полностью. Если нет — блоки удлиняются, простоев не остаётся вовсе, и ответ равен len(tasks).

Одинаковые задачи оказываются в соседних блоках, а длина каждого блока не меньше n + 1, поэтому ограничение кулдауна выполняется.

Решение симуляцией: куча и очередь кулдауна

Формула отвечает только на вопрос «сколько». Если интервьюер просит построить само расписание или условие чуть изменено (например, у задач разная длительность), нужна симуляция.

Жадная стратегия: на каждом такте выполняем доступную задачу, которой осталось больше всего копий. Интуиция та же — самую частую задачу выгодно запускать как можно раньше, иначе в конце она останется одна и будет ждать кулдаун.

Нужны две структуры:

• max-куча с количеством оставшихся копий доступных задач — чтобы быстро брать самую частую (подробнее о куче — в статье про приоритетную очередь);
• очередь кулдауна — задачи, которые недавно выполнялись, с моментом, когда они снова станут доступны. Время у всех растёт одинаково, поэтому обычной FIFO-очереди достаточно: кто раньше вошёл, тот раньше и освободится.

import heapq
from collections import Counter, deque


def least_interval_sim(tasks: list[str], n: int) -> int:
    # в heapq только min-куча, поэтому храним отрицательные количества
    heap = [-c for c in Counter(tasks).values()]
    heapq.heapify(heap)
    cooldown = deque()  # (осталось копий со знаком минус, такт освобождения)
    time = 0

    while heap or cooldown:
        time += 1
        if heap:
            left = heapq.heappop(heap) + 1  # выполнили одну копию
            if left:
                cooldown.append((left, time + n))
        else:
            # доступных задач нет — перематываем простой целиком
            time = cooldown[0][1]
        if cooldown and cooldown[0][1] == time:
            heapq.heappush(heap, cooldown.popleft()[0])

    return time

Пара деталей, на которых легко ошибиться:

• Отрицательные числа. heapq — min-куча. Чтобы доставать максимум, кладём -count. Выполнение задачи — это +1 к отрицательному числу, а ноль означает, что копий не осталось.
• Перемотка простоя. Если куча пуста, не нужно крутить цикл такт за тактом — сразу прыгаем к моменту, когда освободится первая задача в очереди. Это не меняет асимптотику, но выглядит аккуратнее.
• Возврат из кулдауна. Задача, выполненная в такт t, снова доступна в такт t + n + 1. Мы проверяем очередь в конце такта t + n, поэтому на следующей итерации она уже в куче.

Сложность — $O(T \log k)$, где $T$ — длина ответа, а $k$ — число различных задач. Поскольку $k \le 26$, на практике это $O(T)$. Без перемотки простоев симуляция крутит цикл на каждом такте, включая idle; с перемоткой — только на тактах с работой и на прыжках.

Восстановление расписания

Небольшая модификация симуляции возвращает само расписание. В кучу кладём пары (-count, task), а в выходной список — имя задачи или idle:

import heapq
from collections import Counter, deque


def schedule(tasks: list[str], n: int) -> list[str]:
    heap = [(-c, t) for t, c in Counter(tasks).items()]
    heapq.heapify(heap)
    cooldown = deque()  # (осталось копий, задача, такт освобождения)
    result = []

    while heap or cooldown:
        if heap:
            count, task = heapq.heappop(heap)
            result.append(task)
            if count + 1:
                cooldown.append((count + 1, task, len(result) + n))
        else:
            result.append("idle")
        if cooldown and cooldown[0][2] == len(result):
            count, task, _ = cooldown.popleft()
            heapq.heappush(heap, (count, task))

    return result


schedule(list("AAABBB"), 2)
# ['A', 'B', 'idle', 'A', 'B', 'idle', 'A', 'B']

schedule(list("AAAABBBCCD"), 2)
# ['A', 'B', 'C', 'A', 'B', 'C', 'A', 'B', 'D', 'A']

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

Здесь idle вставляется по одному, без перемотки, потому что каждый простой нужно записать в результат.

Вариант с фиксированным порядком: Task Scheduler II

На LeetCode есть задача 2365 «Task Scheduler II», и её любят давать как follow-up. Отличие принципиальное: порядок задач задан и менять его нельзя. Задачи — числа, после выполнения задачи типа t следующая задача того же типа может стартовать не раньше чем через space дней.

Здесь нет никакого выбора, а значит, нет и жадности с кучей. Достаточно идти по задачам и помнить для каждого типа, в какой день он снова доступен:

def task_scheduler_ii(tasks: list[int], space: int) -> int:
    next_free = {}  # тип задачи → первый день, когда её снова можно делать
    day = 0
    for t in tasks:
        day += 1
        if next_free.get(t, 0) > day:
            day = next_free[t]  # ждём, пока кулдаун закончится
        next_free[t] = day + space + 1
    return day


task_scheduler_ii([1, 2, 1, 2, 3, 1], 3)  # 9
task_scheduler_ii([5, 8, 8, 5], 2)  # 6

Разберём первый пример. День 1 — задача 1, она снова доступна с дня 5. День 2 — 2, доступна с дня 6. День 3 хотим сделать 1, но рано — прыгаем на день 5. День 6 — 2, как раз освободилась. День 7 — 3. День 8 — 1: последний раз её делали в день 5, доступна с дня 9 — прыгаем. Итого 9.

Время $O(m)$, память $O(k)$ — хеш-таблица по типам задач. Важно не симулировать ожидание по одному дню: space может быть до $10^9$, и цикл по дням не уложится в лимит. Прыжок day = next_free[t] решает это.

Ещё один родственник — задача «Rearrange String k Distance Apart» (LeetCode 358, премиум): переставить буквы строки так, чтобы одинаковые стояли на расстоянии не меньше k, или сказать, что это невозможно. Это та же симуляция с кучей и очередью, только idle запрещён: если куча опустела, а строка не собрана, ответ — пустая строка. Частный случай k = 2 — открытая задача 767 «Reorganize String».

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

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

• Забыть max с длиной массива. Формула $(f - 1)(n + 1) + k$ без максимума даёт неверный ответ, когда простоев нет: для AAABBBCCDDEE при n = 2 получится 8 вместо 12.
• Считать k неправильно. k — это число различных задач с максимальной частотой, а не сама частота.
• Путать n и n + 1. n — это промежуток между одинаковыми задачами, а длина блока вместе с самой задачей — n + 1.
• Брать задачи из кучи без кулдауна. Если просто вынимать максимум и сразу класть обратно, одна задача будет выполняться подряд. Очередь кулдауна — обязательная часть решения.
• Жадничать не тем. Брать задачи по алфавиту или в порядке появления неоптимально. Пример: AAABC, n = 2. Оптимум — A B C A idle idle A, 7 тактов. Если начать с B и C, получится B C A idle idle A idle idle A — 9.
• Цикл по дням в Task Scheduler II. С большим space это TLE, нужен прыжок.

Как это выглядит на собеседовании

Хорошая последовательность ответа:

1. Сразу заметить, что порядок произвольный, и предложить жадность: самую частую задачу — как можно раньше.
2. Нарисовать каркас A _ _ A _ _ A и вывести формулу. Отдельно проговорить случай, когда простоев нет, и откуда берётся max.
3. Объяснить, почему каркас достижим (раскладка по блокам по кругу).
4. Если попросят расписание или изменят условие — перейти к симуляции с кучей и очередью.

Follow-up вопросы, к которым стоит быть готовым: «а если порядок фиксирован?» (Task Scheduler II), «а если задачи нельзя разделять простоями?» (Reorganize String), «а если у задач разная длительность?» — тогда формула ломается, и остаётся только симуляция.

Задача хорошо стыкуется с другими темами блога: про жадные стратегии в целом — в статье о жадных алгоритмах, про подсчёт частот и кучу — в разборе Top K Frequent Elements, а устройство самой кучи — в статье про двоичную кучу.

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

• Ответ определяется самой частой задачей: $(f_{max} - 1)(n + 1) + k$ тактов — каркас с простоями.
• Ответ не может быть меньше числа задач, поэтому берём max с len(tasks).
• Формула — это и нижняя оценка, и достижимое расписание: остальные задачи раскладываются по блокам по кругу.
• Симуляция: max-куча по оставшимся копиям + FIFO-очередь кулдауна. Нужна, когда просят само расписание.
• При фиксированном порядке выбора нет: хватает хеш-таблицы «тип → день, когда снова доступен» и прыжков через ожидание.