PythonでK個のタスクを完了するための最大時間を見つけるプログラム

問題の概要

各行に3つの値を持つタスクの行列と、もうひとつの値 k が与えられているとします。タスクの中から k 行を選び、それを S と呼ぶことにします。このとき、次の合計値が最小になるように行を選択し、その合計を答えとして返すのが目的です。

max(S[0, 0], S[1, 0], ..., S[k-1, 0]) + max(S[0, 1], S[1, 1], ..., S[k-1, 1]) + max(S[0, 2], S[1, 2], ..., S[k-1, 2])

言い換えると、3つの列それぞれがコストに寄与し、各列のコストは「S 内でのその列の最大値」として計算されます。なお、空のリストに対する最大値は 0 とみなします。

入力例と考え方

tasks = [[2, 3, 3], [4, 5, 2], [4, 2, 3]]、k = 2 の場合、出力は 10 になります。これは 1 行目と 3 行目を選んだケースです。S = [[2,3,3],[4,2,3]] となるため、

• max(S[0,0], S[1,0]) = 4
• max(S[0,1], S[1,1]) = 3
• max(S[0,2], S[1,2]) = 3

となり、合計は 4 + 3 + 3 = 10 となります。

解法のアルゴリズム

この問題を解くには、以下の手順に従います。

1. 関数 util() を定義します。引数として B を受け取ります。
2. リスト B をソートします。
3. yheap を、i が 0 から K-1 の範囲の各 i について -B[i, 1] を要素とするリストとして作成します。
4. yheap をヒープ化(heapify)します。
5. ans := B[K-1, 0] + (-yheap[0]) とします。
6. i が K から B のサイズまでの範囲で、以下を繰り返します。
   • x := B[i, 0]
   • yheap を -B[i, 1] で置き換えます(pushpop 操作)。
   • yheap のサイズが K であることを確認します。
   • y := -yheap[0]
   • ans := ans と x + y のうち小さい方
7. ans を返します。

メインメソッドでは、以下の処理を行います。

1. A が空、または K が 0 の場合は 0 を返します。
2. リスト A をソートします。
3. B を、i が 0 から K-1 の範囲の各 i についてペア [A[i, 1], A[i, 2]] を並べたリストとして作成します。
4. ans := A[K-1, 0] + B 内の各 (y, z) に対する y の最大値 + B 内の各 (y, z) に対する z の最大値
5. i が K から A のサイズまでの範囲で、以下を繰り返します。
   • [A[i][1], A[i][2]] を B に追加します。
   • ans := ans と A[i, 0] + util(B) のうち小さい方
6. ans を返します。

このアルゴリズムでは、1 列目の値を基準に候補を絞り込みながら、残りの 2 列についてはヒープ(優先度付きキュー)を使って上位 K 個の値を効率的に管理することで、全組み合わせを総当たりせずに最小値を求められる点がポイントです。

実装例

import heapq

class Solution:
    def solve(self, A, K):
        if not A or not K:
            return 0

        def util(B):
            B.sort()
            yheap = [-B[i][1] for i in range(K)]
            heapq.heapify(yheap)
            ans = B[K - 1][0] + (-yheap[0])
            for i in range(K, len(B)):
                x = B[i][0]
                heapq.heappushpop(yheap, -B[i][1])
                assert len(yheap) == K
                y = -yheap[0]
                ans = min(ans, x + y)
            return ans

        A.sort()
        B = [[A[i][1], A[i][2]] for i in range(K)]
        ans = A[K - 1][0] + max(y for y, z in B) + max(z for y, z in B)
        for i in range(K, len(A)):
            B.append([A[i][1], A[i][2]])
            ans = min(ans, A[i][0] + util(B))
        return ans

ob = Solution()
tasks = [
    [2, 3, 3],
    [4, 5, 2],
    [4, 2, 3]
]
k = 2
print(ob.solve(tasks, k))

入力

tasks = [
  [2, 3, 3],
  [4, 5, 2],
  [4, 2, 3]
],
k = 2

出力

10