상담원 인원.py

from collections import defaultdict
from itertools import combinations_with_replacement, permutations
import heapq
import math

# 대기열에 n명의 상담원이 있을 때, wait time 총량 계산
def waiting_time(waitings, n):
    total_time = 0
    counsel_list = []
    for _  in range(n):
        heapq.heappush(counsel_list, 0) # 해당 상담원이 가능해질 때까지의 대기시간. 초기엔 0으로 설정
    
    for (start, end) in waitings:
        prev_end = heapq.heappop(counsel_list)
        if start > prev_end: # 완료된 이후에 찾아오면 대기시간 없음
            heapq.heappush(counsel_list, end)
        else:
            wait_time = prev_end - start # 끝나는 시간보다 먼저 찾아온 만큼 대기
            total_time += wait_time
            heapq.heappush(counsel_list, wait_time + end) # 대기한 시간만큼 종료가 길어진다.
    
    return total_time
    
    

def solution(k, n, reqs):
    # 모든 경우의 수를 뽑아보자.
    # 문제 제약에 따르면 각 유형에 최소 1명은 필요하다 했으니, 유형 하나에 할당할 수 있는 최댓값은 n - (k - 1) (5개면, 최대한으로 부여할 1개 뺀 나머지 (= k-1) 에는 1씩 할당.)
    cases = set()
    for c in combinations_with_replacement([i for i in range(1, n -k + 1 + 1)], k):
        if sum(c) == n:
            # 이 숫자조합을 상담에 맞게 변환할 수 있다. (1,1,2) -> (2,1,1) (1,2,1), (1,1,2)로 다시 조합 가능
            case = set(permutations(c, k))
            cases = cases | case
    
    # 시작시간, 끝나는 시간으로 변환    
    waiting_lists = defaultdict(list)
    for req in reqs:
        start, duration, types = req
        # 특정 상담유형 'types' 에 누가 언제 들어와서 언제 끝나는지
        waiting_lists[types-1].append((start, start + duration))
    
    result = math.inf
    
    for case in cases:
        wait_time = 0
        for i in range(k):
            # i번 상담유형에 case[i] 만큼의 상담원이 들어갔을 때의 대기시간
            wait_time += waiting_time(waiting_lists[i], case[i])
        result = min(wait_time, result)
    
    return result