Python에서 K개 작업을 완료하는 데 드는 최소 시간을 찾는 프로그램

문제 설명

각 행에 3개의 값이 담긴 작업(task) 행렬이 하나 있고, 또 다른 값 k가 주어진다고 가정해 봅시다. 우리는 tasks에서 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이 됩니다. 첫 번째 행과 마지막 행을 선택하면 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이 됩니다.

풀이 접근 방법

이 문제는 정렬과 최소 힙(min-heap)을 함께 활용하면 효율적으로 해결할 수 있습니다. 단계별로 살펴보겠습니다.

util() 함수

util() 함수를 정의합니다. 이 함수는 리스트 B를 인자로 받습니다.
B를 정렬합니다.
yheap := 범위 0 ~ K-1의 각 i에 대해 -B[i, 1] 값을 담은 리스트로 초기화합니다.
yheap을 힙 구조로 만듭니다(heapify).
ans := B[K - 1, 0] + (-yheap[0])으로 초기화합니다.
i를 K부터 B의 크기까지 반복합니다.
- x := B[i, 0]
- yheap에 -B[i, 1]을 삽입하되 힙의 크기를 K로 유지합니다(heappushpop).
- y := -yheap[0]
- ans := ans와 x + y 중 더 작은 값
ans를 반환합니다.

메인 로직

A가 비어 있거나 K가 0이면 0을 반환합니다.
A를 정렬합니다.
B := 범위 0 ~ K-1의 각 i에 대해 [A[i, 1], A[i, 2]] 쌍으로 이루어진 리스트를 만듭니다.
ans := A[K - 1, 0] + B에 있는 y들의 최댓값 + B에 있는 z들의 최댓값으로 초기화합니다.
i를 K부터 A의 크기까지 반복합니다.
- B에 [A[i][1], A[i][2]]를 추가합니다.
- ans := ans와 A[i, 0] + util(B) 중 더 작은 값
ans를 반환합니다.

동작 원리 요약

핵심 아이디어는 다음과 같습니다. 먼저 작업을 첫 번째 열 값을 기준으로 오름차순 정렬하면, 특정 행을 기준점으로 삼았을 때 그 행의 첫 번째 열 값이 선택된 행들 중 최댓값이 됩니다. 이후 나머지 두 열에 대해서는 힙을 활용해 후보군 안에서 가장 작은 '최댓값의 합' 조합을 효율적으로 탐색합니다. 이를 통해 모든 조합을 일일이 확인하는 브루트포스 방식보다 훨씬 빠르게 최적해를 구할 수 있습니다.

구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

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