NeetCode 150 [Heap Priority Queue]medium:Task Scheduler

NeetCodeのSolutionを書いていく

Task Scheduler

問題

CPUタスクを表すtasksと整数nが与えられます。tasks[i]にはAからZの大文字の英語が入ります。
それぞれのCPUサイクルで一つのタスクを完了でき、どのタスクから完了してもよいです。
唯一の制限は、CPUをクールダウンするために、同一のタスクは少なくともnサイクル開ける必要があることです。
すべてのタスクを完了するために必要な最小のCPUサイクルを返しなさい。

例

Input: tasks = ["X","X","Y","Y"], n = 2
Output: 5

Input: tasks = ["A","A","A","B","C"], n = 3
Output: 9

制約

• 1 <= tasks.length <= 1000
• 0 <= n <= 100
• 計算量：O(m) mは配列数
• メモリ：O(1)

メモ

「同一のタスク」の場合は間をn空けないといけない。
つまり「同一ではないタスク」であれば連続で実行できる。
初期状態で同一のタスクがいくつあるかが問題なのではなく、
同一ではないタスク同士を優先して並べて行った時に最後に残った同一タスクをどう扱うかが問題。
数が多い順に処理をしながら、一旦pending_queueに退避、復帰する流れ

import heapq
from collections import Counter, deque
from typing import List, Tuple


class Solution:
    def leastInterval(self, tasks: List[str], n: int) -> int:
        # key毎の数をカウント
        key_count = Counter(tasks)

        # max_heapとして使いたいのでカウントの値を逆にする
        inputs = [-v for v in key_count.values()]

        # heapifyで順番を更新
        heapq.heapify(inputs)

        # 退避キューを作成
        invalid_queue: deque[Tuple[int, int]] = deque()  # [-cnt, idle_time]のペア

        # 処理回数を初期化
        time = 0

        # 入力か、退避キューに値があれば処理をする
        while inputs or invalid_queue:
            time += 1

            # inputsがない場合invalid_queueの先頭を確認
            # 取り出してよいがイミングであれば取り出す
            if not inputs:
                time = invalid_queue[0][1]
            else:
                # inputsから値取り出し
                cnt = 1 + heapq.heappop(inputs)
                if cnt:
                    # 現在の残りの数と次回取り出して良いタイミングをpush
                    invalid_queue.append((cnt, time + n))

            # キューの先頭の値が取り出して良いタイミングならinputsに移す
            if invalid_queue and invalid_queue[0][1] == time:
                heapq.heappush(inputs, invalid_queue.popleft()[0])

        return time