Start learning free

LESSON 391 OF 540 · HEAP

Task Scheduler in Scala: Max-Heap and 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 <= 10⁵
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 15

Guided

Ctrl + Enter to check

SCALA · HEAP

Task Scheduler in Scala, explained

Problem overview

The method leastInterval receives tasks, an Array[Char] of uppercase letters, and a cooldown n. 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. Two copies of the same task must be at least n units apart: after running a task, the CPU needs n other units (other tasks or idle) before it can run that task again. Tasks can run in any order. The method returns the fewest units of time needed to run them all.

For tasks = [A, A, A, B, B, B] and n = 2, one best order is A B idle A B idle A B: 8 units. With only two kinds of task, the CPU has to idle twice. With n = 0 there is no cooldown, any order works, and the answer is 6.

Approach: always run the task with the most copies left

Only how many copies each task has matters, not the order of the input, so the first step is to count them. There are at most 26 letters.

Then simulate the CPU. In each unit, run the task with the most copies left. Those copies need the most room between them, so starting them early leaves the fewest idle units at the end. After a task runs, it waits out its cooldown before it can be picked again.

Two structures do the work:

a max-heap of the counts of tasks that are ready to run, with the biggest count on top (Scala's PriorityQueue is a max-heap);
a queue of tasks that are cooling down, each stored as a tuple (left, readyTime): the copies it has left and the time at whose end its cooldown is over, so it can run again one unit later.

Tasks join the queue in the order they ran, so they also leave it in order of ready time, and the front is always the first one ready. When the heap is empty but tasks are still cooling, the CPU idles, and the time jumps straight to the moment the first one is ready instead of counting each idle unit.

The heap never holds more than 26 counts, so each heap operation is effectively constant: the total is O(T log 26) for T tasks, using O(1) extra space (at most 26 counts).

The statement's note also mentions a counting formula for this problem. This lesson practises the simulation.

Step by step

Count the copies: val count = new Array[Int](26), one place per letter, starting at zeros.
Loop over the tasks with for (task <- tasks).
Add one to the letter's count at count(task - 'A'). Subtracting 'A' turns A into 0, B into 1, up to Z as 25.
Create the heap: PriorityQueue[Int](), whose top is the biggest count.
Loop over the counts with for (c <- count).
Enqueue the nonzero ones: letters that never appear have nothing to run.
Create the cooling queue: mutable.Queue[(Int, Int)]() of (left, readyTime) tuples.
Start time at 0, a var. It is the answer once everything has run.
Simulate while work remains: heap.nonEmpty || cooling.nonEmpty. The work is done only when nothing waits to run and nothing is cooling.
Start a new unit: time += 1.
If the heap is empty, no task can run, so the CPU idles until the first cooling task is ready: time = cooling.head._2. This counts all the idle units at once.
Otherwise run a task: val left = heap.dequeue() - 1 is how many copies the most frequent task has after this one.
Start its cooldown if left > 0: enqueue (left, time + n) into cooling. A task run at time can run again at time + n + 1. A task with no copies left is done.
Release a ready task: if cooling.nonEmpty && cooling.head._2 == time, take it out with cooling.dequeue() and put its count, ._1, back into the heap.
Return time, the unit in which the last task ran.

The solution in Scala

In a file of your own, the code needs import scala.collection.mutable and import scala.collection.mutable.PriorityQueue at the top.

class Solution {
    def leastInterval(tasks: Array[Char], n: Int): Int = {
        val count = new Array[Int](26)
        for (task <- tasks) {
            count(task - 'A') += 1
        }
        val heap = PriorityQueue[Int]()
        for (c <- count) {
            if (c > 0) {
                heap.enqueue(c)
            }
        }
        val cooling = mutable.Queue[(Int, Int)]()
        var time = 0
        while (heap.nonEmpty || cooling.nonEmpty) {
            time += 1
            if (heap.isEmpty) {
                time = cooling.head._2
            } else {
                val left = heap.dequeue() - 1
                if (left > 0) {
                    cooling.enqueue((left, time + n))
                }
            }
            if (cooling.nonEmpty && cooling.head._2 == time) {
                heap.enqueue(cooling.dequeue()._1)
            }
        }
        time
    }
}

For tasks = [A, A, A, B, B, B] and n = 2, the loop goes like this:

time=1 ran task, left=2
time=2 ran task, left=2
idle jump to 3
  ready: (2,3)
time=4 ran task, left=1
  ready: (2,4)
time=5 ran task, left=1
idle jump to 6
  ready: (1,6)
time=7 ran task, left=0
  ready: (1,7)
time=8 ran task, left=0

The two tasks each have 3 copies. Units 1 and 2 run one copy of each, and both go into the cooling queue with 2 copies left, ready at the end of units 3 and 4. The heap is now empty, so the time jumps to 3, the first ready time: that is the idle unit. The first task returns to the heap and runs in unit 4, the second returns at the end of unit 4 and runs in unit 5. Another idle jump covers unit 6, and the last two copies run in units 7 and 8. The result is 8.

The Scala details worth a closer look:

Characters are numbers. 'C' - 'A' is the Int 2, so task - 'A' works as an index. count(task - 'A') += 1 updates the array element in place.
Tuples and ._1, ._2. A tuple's items are read as ._1 and ._2, counting from 1. cooling.head._2 is the ready time of the front task, and cooling.dequeue()._1 is the number of copies it has left.
enqueue((left, time + n)) has two pairs of parentheses: the outer pair is the call, the inner pair builds the tuple.
Max-heap of Int. PriorityQueue[Int]() returns the biggest element from dequeue(), so no custom ordering is needed.
head and dequeue() on the queue. head looks at the front without removing it; dequeue() removes and returns it. Both throw an exception on an empty queue, which is why the code checks cooling.nonEmpty first and only reads head in the idle branch when the loop condition guarantees something is cooling.
Heap and queue don't share a type here. The heap holds Int counts, the queue holds tuples; the mutable.Queue is first in, first out, which is what keeps the cooling tasks in ready-time order.

Complexity

Time: O(T log 26), where T is the number of tasks. Counting takes O(T). Each unit of simulated work runs one task, with a few heap and queue operations; the heap holds at most 26 counts. Idle stretches are skipped in one jump, so the loop doesn't depend on how large n is.

Space: O(1). The count array has 26 places, and the heap and the queue together never hold more than 26 tasks.

Edge cases and common mistakes

No cooldown. With n = 0, time + n equals time, so the task is ready again straight away, and AAABBB takes 6 units.
One kind of task. AA with n = 3 gives 5: A, three idle units, A. A single A with n = 100 gives 1, because the last copy needs no cooldown after it.
Plenty of other tasks. For ABCD with n = 5, each task runs once, nothing repeats, and the answer is 4 with no idling.
Many different tasks filling the gaps. Six A's and six other letters with n = 2 give 16: the other tasks fill the first gaps and the CPU idles in the rest.
No idle time at all. AAABBBCC with n = 2 gives 8, the number of tasks, because there is always something else to run.
Counting the cooldown as a gap of n units between starts. Two copies of a task must have at least n other units between them, so the next start is n + 1 units later, not n.

Testing

AAABBB with n = 2 gives 8.
AAABBB with n = 0 gives 6.
AAAAAABCDEFG with n = 2 gives 16.
A with n = 100 gives 1.
AA with n = 3 gives 5.
ABCD with n = 5 gives 4.
AAABBBCC with n = 2 gives 8.

Keep going with Scala500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free