Read-only web browser

Solutions

Browse the interview solutions folder with syntax highlighting.

task-scheduler.py

heap/task-scheduler.py · Python · 3.2 KiB · 2026-03-25 20:59

Back to folder
# Always process the most frequent task first (max heap), use a queue to store the tasks that are cooling down and when they can be executed again add them back to the heap.
# heap

# O(n log k) time, O(k) space
# but since k can be at most 26 (the number of uppercase letters), we can consider it O(n) time, O(1) space
def leastInterval(tasks: list[str], n: int) -> int:
    # Collect frequencies of tasks
    from collections import defaultdict
    freq = defaultdict(int)
    for t in tasks:
        freq[t] += 1

    # Put the frequencies in a max heap
    # INTUITION: we process first the most frequent one
    import heapq
    maxh = []
    for task, f in freq.items():
        heapq.heappush(maxh, -f) # NB: no need to store the task id

    # Queue to store task along with the next time they become executable
    from collections import deque
    q = deque() # (count, next time of execution)

    time = 0
    while q or maxh:
        time += 1

        if maxh:
            # If the heap is not empty, process the most frequent task
            f = heapq.heappop(maxh)

            # Decrement the frequency
            # and if it's still positive, add it back to the queue
            f += 1
            if f != 0:
                q.append((f, time + n))
        
        # If current time is the same as execution time of the task at the head of the queue,
        # add it back to the max heap for execution
        if q and time == q[0][1]:
            cnt, _ex_time = q.popleft()
            heapq.heappush(maxh, cnt)

    return time


# Greedy + counting - O(n) time, O(1) space
def leastInterval(tasks: list[str], n: int) -> int:
    count = [0] * 26
    for task in tasks:
        count[ord(task) - ord('A')] += 1

    # Get the idle times we need to insert between the most frequent task
    # if count['A']=5 and n=2: A _ _ A _ _ A _ _ A _ _ A => (5-1) * 2 = 8 idle times
    count.sort()
    maxf = count[25]
    idle = (maxf - 1) * n

    # Fill the idle times with the remaining tasks, starting from the most frequent one
    for i in range(24, -1, -1):
        idle -= min(maxf - 1, count[i])

    # The total time is the number of tasks plus the idle times we need to insert,
    # but if idle is negative it means that we can fill all the idle times with the remaining tasks, so we just return the number of tasks
    return max(0, idle) + len(tasks)


# Brute-force - O(n * t) time, O(t) space
def leastInterval(tasks: list[str], n: int) -> int:
    count = [0] * 26
    for task in tasks:
        count[ord(task) - ord('A')] += 1

    arr = []
    for i in range(26):
        if count[i] > 0:
            arr.append([count[i], i])

    time = 0
    processed = []
    while arr:
        maxi = -1
        for i in range(len(arr)):
            # Check if the task was not executed in the cooling down period
            if all(processed[j] != arr[i][1] for j in range(max(0, time - n), time)):
                if maxi == -1 or arr[maxi][0] < arr[i][0]:
                    maxi = i

        time += 1
        cur = -1
        if maxi != -1:
            cur = arr[maxi][1]
            arr[maxi][0] -= 1
            if arr[maxi][0] == 0:
                arr.pop(maxi)
        processed.append(cur)
    return time