Task Scheduler With Cooldown: Max-Heap Simulation vs the Formula

The task scheduler problem gives you a list of tasks, a cooldown, and one CPU, and asks for the shortest schedule. It has two correct answers that interviewers like for different reasons: a max-heap simulation that you can explain step by step and that produces the schedule itself, and a counting formula that fits on one line but needs an argument to be believed. This article builds both, shows they agree, and checks them against an exhaustive search.

The problem it solves

Each task takes one second. Identical tasks must be at least n seconds apart: after running a at second t, the next a can run at t + n + 1 at the earliest. When nothing is allowed to run, the CPU idles, and idle seconds count. Tasks may run in any order. Return the minimum total time.

For a a a b b c with n = 2 the answer is 7, for example a b c a b _ a: the three as need two gaps of two seconds each, and there are not enough other tasks to fill both gaps completely.

The intuition

The task with the most runs left is the bottleneck. Every run of a forces n seconds before the next one, so the task with the most copies determines how long the schedule must be stretched. Running it as early and as often as it is allowed keeps its chain of forced gaps as short as possible, and the other tasks fill those gaps. Leaving the busiest task for later only packs its gaps at the end, where nothing is left to fill them.

That gives the simulation. Keep a max-heap of remaining counts for tasks that are allowed to run now, and a queue of tasks that are cooling down, each stamped with the second it becomes legal again. Every second: move any task whose cooldown has ended back into the heap; if the heap is not empty, run its top and, if copies remain, park it in the queue; otherwise idle. The clock advances either way.

The formula comes from the same picture. Let f be the highest count and m the number of tasks that have it. Lay the busiest task out in f - 1 frames of length n + 1, then one final frame with just the m tied tasks:

(f - 1) × (n + 1) + m

Everything else fits into the frames' empty slots. If there are more tasks than slots, there are no idles at all and the answer is simply the number of tasks. So the answer is the larger of the two.

Watch it run

The animation schedules three a, two b and one c with a cooldown of 2. At t = 1 it runs a, the task with the most runs left, and parks it until t = 4. At t = 2 it runs b and parks it until 5. At t = 3 it runs c, its last run. At t = 4 a is ready again and runs, parked until 7. At t = 5 b runs for the last time. At t = 6 every remaining task is still cooling, so the second is spent idling, and it still counts. At t = 7 the last a runs: six tasks, 7 seconds. With enough distinct tasks there would have been no gap at all.

The code

The simulation from the lesson, extended to return the schedule it builds, with _ for an idle second:

import heapq
from collections import Counter, deque

def schedule(tasks, n):
    heap = [(-count, task) for task, count in Counter(tasks).items()]
    heapq.heapify(heap)                         # most runs left on top
    cooling = deque()                           # (ready_at, -count, task), in time order
    time, order = 0, []
    while heap or cooling:
        time += 1
        if cooling and cooling[0][0] <= time:
            _, neg, task = cooling.popleft()
            heapq.heappush(heap, (neg, task))   # its cooldown is over
        if heap:
            neg, task = heapq.heappop(heap)
            order.append(task)
            if neg + 1 < 0:
                cooling.append((time + n + 1, neg + 1, task))
        else:
            order.append("_")                   # nothing is allowed: idle, and it counts
    return time, "".join(order)

print(schedule("aaabbc", 2))                    # (7, 'abcab_a')
print(schedule("abcd", 2))                      # (4, 'abcd')
print(schedule("aaaa", 3))                      # (13, 'a___a___a___a')

The formula, with the two cases that matter: idles forced by the busiest task, and enough variety to need none:

def least_interval(tasks, n):
    counts = Counter(tasks)
    f = max(counts.values())
    m = sum(1 for c in counts.values() if c == f)   # tasks tied for the top count
    return max(len(tasks), (f - 1) * (n + 1) + m)

print(least_interval("aaabbc", 2))              # 7
print(least_interval("aaabbb", 2))              # 8
print(least_interval("aaabbbcccdd", 2))         # 11

Both against an exhaustive search that tries every legal choice each second, including idling on purpose, on 400 small random inputs:

import random
from functools import lru_cache

def brute(tasks, n):
    names = sorted(set(tasks))
    start = tuple(tasks.count(t) for t in names)

    @lru_cache(maxsize=None)
    def best(left, wait):                       # wait[i]: seconds until task i is legal
        if not any(left):
            return 0
        tick = tuple(max(0, w - 1) for w in wait)
        options = [1 + best(left, tick)] if any(wait) else []   # idle, when it changes anything
        for i, c in enumerate(left):
            if c and wait[i] == 0:
                nl = left[:i] + (c - 1,) + left[i + 1:]
                nw = tick[:i] + (n,) + tick[i + 1:]
                options.append(1 + best(nl, nw))
        return min(options)

    return best(start, (0,) * len(names))

def legal(order, n):
    last = {}
    for t, task in enumerate(order):
        if task != "_":
            if task in last and t - last[task] <= n:
                return False
            last[task] = t
    return True

random.seed(17)
ok = True
for _ in range(400):
    tasks = "".join(random.choice("abc") for _ in range(random.randint(1, 7)))
    n = random.randint(0, 3)
    time, order = schedule(tasks, n)
    ok &= time == least_interval(tasks, n) == brute(tasks, n) == len(order)
    ok &= legal(order, n) and sorted(order.replace("_", "")) == sorted(tasks)
print(ok)                                       # True

The complexity

With T tasks, k distinct task names and a total time of L seconds:

• Simulation: O(L log k), one heap operation per second. L can exceed T because of idles: aaaa with n = 3 takes 13 seconds for 4 tasks. The heap and the queue hold at most k entries.
• Formula: O(T) to count, O(k) extra space, and no loop over seconds at all.
• Exhaustive search: exponential in the number of task kinds; only useful as a test oracle.

Where it goes wrong

• Not counting idle seconds. The answer is total time, not the number of tasks.
• Waking a task one second early. A task run at t is legal again at t + n + 1, not t + n.
• Forgetting ties in the formula. With aaabbb and n = 2 the last frame holds both a and b, so m = 2 and the answer is 8, not 7.
• Dropping the max with the task count. When many distinct tasks fill every gap, (f - 1) × (n + 1) + m underestimates; aaabbbcccdd with n = 2 needs 11 seconds, one per task.
• Picking the first allowed task alphabetically. On abbb with n = 1 that runs ab_b_b in 6 seconds, where running the busiest task first gives bab_b in 5.

When it shows up in interviews

It is a common medium in the heap and greedy sections. Some interviewers want the formula and a proof; others want the simulation because the follow-ups need it: "print the schedule", "tasks arrive with priorities", or "the order is fixed and you can only wait", which changes the problem into a single pass with a map of the next legal time per task. The same most-frequent-first reasoning solves reorganize a string, which is this problem with n = 1 and no idling allowed.

How to say it in an interview

"The task with the most copies is the bottleneck, so each second I run the allowed task with the most runs left. I keep a max-heap of counts for tasks that are ready and a queue of tasks cooling down, stamped with the second they become legal again, and the clock advances even when I have to idle. That is O(L log k). There is also a closed form: with top count f shared by m tasks, the answer is the max of the number of tasks and (f - 1)(n + 1) + m."

The same greedy on a tighter rule is in reorganize a string, and choosing the best item repeatedly from a heap is covered in top k elements: heap vs sort.