Pythonで同種タスク間にk時間の間隔を空けて全タスクを完了するのに必要な最小時間を求めるプログラム

問題の概要

整数のリスト tasks（各要素はタスクの種類を表します）と、負でない整数 k が与えられます。各タスクの完了には1単位の時間がかかり、タスクは必ず与えられた順序どおりに実行しなければなりません。ただし、同じ種類のタスクを2回実行する間には、少なくとも k 単位の時間を空ける必要があります。任意の時点で選択できるのは「タスクを実行する」か「待機する」かのいずれかです。すべてのタスクを完了するまでに必要な合計時間を求めてください。

入力例

たとえば、tasks = [0, 1, 1, 2]、k = 2 の場合、出力は 6 になります。

最初の2つのタスクは種類が異なるため、間隔を空けずに連続して実行できます。しかし時刻2で次のタスクは1番目と同じ種類なので、2タイムスロット待機してから実行し、最後に残りの別種類のタスク（タイプ2）を実行します。つまり実行の流れは次のようになります。

[0, 1, 待機, 待機, 1, 2]

このことから、必要なタイムスロットの総数は6であることがわかります。

解法のアプローチ

この問題は、タスクの種類ごとに「次に実行可能になる時刻」を記録しておくことで、効率よく解くことができます。手順は以下のとおりです。

• 現在時刻を表す tick を0で初期化します。
• タスク種類ごとの次回実行可能時刻を格納するマップ slot を用意します。
• リスト内の各タスク t に対して、以下を繰り返します。
  • t が slot に登録されていれば、その値を tf とします。
  • tf が存在し、かつ tf - tick > 0 の場合は、tick を tf まで進めます（クールダウン時間ぶん待機）。
  • tick を1増やしてタスクを実行します。
  • slot[t] = tick + k として、この種類のタスクを次に実行できる最早時刻を更新します。
• 最後に tick を返します。

Pythonでの実装例

理解を深めるために、以下の実装例を見てみましょう。

def solve(tasks, k):
    tick = 0
    slot = {}
    for t in tasks:
        tf = slot.get(t)
        if tf is not None and tf - tick > 0:
            tick += tf - tick
        tick += 1
        slot[t] = tick + k

    return tick

tasks = [0, 1, 1, 2]
k = 2
print(solve(tasks, k))

入力

[0, 1, 1, 2]

出力

6

計算量

タスクの総数を n とすると、リストを一度走査するだけで処理が完了するため、時間計算量は O(n) です。また、登場するタスクの種類数を m とすると、マップに保存する情報は種類ごとに1件のみなので、空間計算量は O(m) となります。