Start learning free

LESSON 391 OF 540 · HEAP

Task Scheduler in Go: O(T log 26) with a Max-Heap and a Cooldown Queue

Practice O(T log 26) Read the full explanation

A CPU has to run a list of tasks. Each task is an uppercase letter from A to Z; equal letters are copies of the same task. Running one task takes one unit of time, and in each unit the CPU either runs one task or stays idle.

There is a cooldown: two copies of the same task must be at least n units apart, so after running a task the CPU needs n other units (other tasks, or idle) before it can run that task again. The tasks can run in any order.

Return the fewest units of time the CPU needs to run all the tasks.

 

Example 1:

Input: tasks = ["A","A","A","B","B","B"], n = 2
Output: 8
Explanation: One best order is A B idle A B idle A B.
After an A, two other units must pass before the next A, and B has the same rule. With only two kinds of task, the CPU has to idle twice.

Example 2:

Input: tasks = ["A","A","A","B","B","B"], n = 0
Output: 6
Explanation: With no cooldown, any order works, for example A A A B B B, and no unit is idle.

Example 3:

Input: tasks = ["A","A","A","A","A","A","B","C","D","E","F","G"], n = 2
Output: 16
Explanation: One best order is A B C A D E A F G A idle idle A idle idle A.
The six A's need two units between each other; the other tasks fill the first gaps, and the CPU idles in the rest.

 

Constraints:

  • 1 <= tasks.length <= 105
  • Every task is an uppercase English letter.
  • 0 <= n <= 100

Note:

Your algorithm should run in O(T log 26) time, where T is the number of tasks (log 26 because a heap here never holds more than 26 kinds of task), and use O(1) extra space: at most 26 counts.

The tutorial simulates the CPU with a max-heap: in each unit, run the task with the most copies left, then let it wait in a queue until its cooldown is over. When no task is ready, the time jumps straight to the moment the first waiting task is ready again. (The answer also has a counting shortcut: if the most frequent task appears m times and c tasks appear that often, it is the bigger of tasks.length and (m - 1) · (n + 1) + c.)

STEP 1 OF 16
    Ctrl + Enter to check

    GO · HEAP

    Task Scheduler in Go, explained

    Problem overview

    A CPU has to run a list of tasks, given to leastInterval(tasks []byte, n int) int. Each task is an uppercase letter, and equal letters are copies of the same task. Running a task takes one unit of time, and in each unit the CPU either runs one task or stays idle.

    There is a cooldown: two copies of the same task must be at least n units apart. After running a task, the CPU needs n other units, either other tasks or idle time, before it can run that task again. The tasks can run in any order, and the function returns the fewest units of time needed to run them all.

    For tasks = A A A B B B and n = 2 the answer is 8: A B idle A B idle A B. With only two kinds of task, the CPU has to idle twice. With n = 0 there is no cooldown and the answer is 6, the number of tasks.

    The lesson asks for O(T log 26) time, where T is the number of tasks and the 26 is the number of different letters, and for O(1) extra space: at most 26 counts.

    Approach: always run the task with the most copies left

    The names of the tasks don't matter, only how many copies each one has. So the first step is to count them: 26 counters, one per letter.

    Then the CPU is simulated. In each unit, run the task that has the most copies left among the ones that are not cooling down. A task with many copies left needs the most room between its copies, so starting it early leaves the fewest idle units at the end. A max-heap of counts gives that task in O(log 26).

    After running a task, it has to wait. The simulation keeps a queue, cooling, of the tasks that are waiting. Each entry holds {left, readyTime}: the copies the task has left, and the time unit at whose end its cooldown is over. A task run at time t can run again at t + n + 1, after n other units, so its ready time is t + n. Tasks join the queue in the order they run, so they also leave it in order of ready time: only the front of the queue can be ready first.

    Sometimes the heap is empty while tasks still cool down. The CPU has nothing to run, so it idles until the front of the queue is ready. Instead of stepping through each idle unit, the time jumps straight to that ready time, which counts all the idle units at once.

    Here is the simulation for A A A B B B with n = 2, printing one line per turn of the loop:

    time 1 run, left 2
    time 2 run, left 2
    idle, jump to 3
    time 4 run, left 1
    time 5 run, left 1
    idle, jump to 6
    time 7 run, left 0
    time 8 run, left 0
    total 8

    Units 1 and 2 run the two tasks, the heap is then empty, and unit 3 is idle. The same pattern repeats, and the last pair runs at units 7 and 8. For A A A B with n = 3 the simulation gives 9: A runs at units 1, 5 and 9, with B in unit 2 and the rest idle.

    The statement also mentions a counting shortcut: if the most frequent task appears m times and c tasks appear that often, the answer is the bigger of len(tasks) and (m - 1) * (n + 1) + c. This lesson practises the heap simulation, which builds the schedule instead of only counting it.

    Step by step

    1. Count the copies. count := make([]int, 26) has one place per letter: index 0 for A up to 25 for Z.
    2. Loop over the tasks with for _, task := range tasks.
    3. Add one to the right counter: count[task-'A']++. Letters are numbers in alphabetical order, so subtracting 'A' turns A into 0, B into 1, and so on.
    4. Create the heap: taskHeap := &maxHeap{}. The lesson provides the maxHeap type, where the biggest count is on top. The variable can't be called heap, because that is the package name.
    5. Loop over the counts with for _, c := range count.
    6. Push only the non-zero counts: if c > 0 { heap.Push(taskHeap, c) }. Letters that never appear have nothing to run.
    7. Create the cooldown queue: cooling := [][]int{}, with entries {left, readyTime}.
    8. Start the clock: time := 0. It counts the units used so far and is the answer at the end.
    9. Loop while work remains: for taskHeap.Len() > 0 || len(cooling) > 0. The work is done only when nothing is waiting to run and nothing is cooling down.
    10. Start a new unit: time++.
    11. Nothing can run? If taskHeap.Len() == 0, jump the clock: time = cooling[0][1]. This makes the idle units count in one step.
    12. Otherwise run a task: in the else branch, left := heap.Pop(taskHeap).(int) - 1 takes the biggest count and uses up one copy. heap.Pop returns any, and .(int) turns it back into an int.
    13. Queue it if copies remain: if left > 0, append []int{left, time + n} to cooling and store the grown slice back. A finished task is not queued.
    14. Release a ready task. Still inside the loop, check len(cooling) > 0 && cooling[0][1] == time. The queue is in order of ready time, so only its front can be ready, and the length check comes first so the front is only read when it exists.
    15. Move it back to the heap: push cooling[0][0] with heap.Push, then drop it from the queue with cooling = cooling[1:].
    16. Return time. It is the unit in which the last task ran.

    The solution in Go

    func leastInterval(tasks []byte, n int) int {
        count := make([]int, 26)
        for _, task := range tasks {
            count[task-'A']++
        }
        taskHeap := &maxHeap{}
        for _, c := range count {
            if c > 0 {
                heap.Push(taskHeap, c)
            }
        }
        cooling := [][]int{}
        time := 0
        for taskHeap.Len() > 0 || len(cooling) > 0 {
            time++
            if taskHeap.Len() == 0 {
                time = cooling[0][1]
            } else {
                left := heap.Pop(taskHeap).(int) - 1
                if left > 0 {
                    cooling = append(cooling, []int{left, time + n})
                }
            }
            if len(cooling) > 0 && cooling[0][1] == time {
                heap.Push(taskHeap, cooling[0][0])
                cooling = cooling[1:]
            }
        }
        return time
    }

    The Go details worth a closer look:

    • []byte and task-'A'. Each task is a byte, and 'A' is an untyped constant that converts to byte here. The difference task-'A' is a byte, so for 'C' it is 2, and any integer type is valid as a slice index. The lesson restricts tasks to A to Z, so the index is always between 0 and 25.
    • count[...]++. Go allows ++ on an indexed element, but only as a statement, not as an expression.
    • The heap holds ints. The type is provided for container/heap, so values go in with heap.Push and come out with heap.Pop. Because Pop returns any, .(int) is a type assertion. It would panic if the value weren't an int, which can't happen here.
    • Queue with a slice. cooling = cooling[1:] drops the first element without copying. With at most 26 tasks waiting at once, a plain slice is enough for a queue.
    • The name time. It is an ordinary variable here. It would only clash with the standard time package in a file that imports it.
    • Why the jump lands right. After time = cooling[0][1], the next if sees cooling[0][1] == time and puts that task back in the heap. The units skipped over are the idle ones, and the clock now shows the last idle unit: no task ran in it, and the next turn starts the unit after it.

    Complexity

    Time: O(T log 26). Counting is O(T). Every copy of every task is popped from the heap once and pushed back at most once more, and the heap holds at most 26 counts, so each operation costs O(log 26). Idle stretches are skipped in one step, so every turn of the loop either runs a task or jumps to a moment when a task is ready and then runs it. The number of turns is at most about twice the number of tasks, whatever the idle time. With 100,000 copies of A and n = 100, the function returns 10099900.

    Space: O(1). The counts hold 26 integers, the heap at most 26, and the queue at most 26 entries.

    Edge cases and common mistakes

    • No cooldown. With n = 0, A A A B B B gives 6. A task is queued with ready time equal to the current time, and the check right after releases it in the same turn.
    • One task. A with n = 5 gives 1: there is no second copy, so nothing has to wait.
    • A single repeated task. A A A with n = 2 gives 7: A, two idle units, A, two idle units, A.
    • Plenty of other tasks. A B C D E F with n = 3 gives 6. Every task is different, so nothing ever waits.
    • Many tasks, big cooldown. A A B with n = 100 gives 102: A, B, 99 idle units, A.
    • Several tasks with the same count. A A B B C C with n = 1 gives 6, with no idle unit: the three tasks take turns.
    • Reading the queue front without checking. cooling[0] panics on an empty slice. Both reads of the front are guarded: the jump runs only while some task is cooling down, and the release check tests len(cooling) > 0 first.
    • Running the smallest count first. That leaves the biggest task for last, with no other task to fill its cooldown, and the schedule gets longer.
    • Off by one in the cooldown. A task run at time t needs n units between its copies, so it may run again at t + n + 1. The ready time t + n is the unit before that.
    • Simulating every idle unit. It works, but it spends time on the idle stretches. The jump avoids that.

    Testing

    • tasks = A A A B B B, n = 2 gives 8.
    • tasks = A A A B B B, n = 0 gives 6.
    • tasks = A A A A A A B C D E F G, n = 2 gives 16.
    • tasks = A, n = 5 gives 1.
    • tasks = A A A, n = 2 gives 7.
    • tasks = A A A B, n = 3 gives 9.
    • tasks = A A B, n = 100 gives 102.
    • tasks = A B C D E F, n = 3 gives 6.
    • tasks = A A B B C C, n = 1 gives 6.
    Keep going with Go500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free