[프로그래머스] 상담원 인원 (Python)

먼저, 특정 유형에 할당된 멘토가 m명일 때 최소 대기시간을 구하는 함수를 작성해 보자.

import heapq
from collections import deque

def get_standby_time(arr, m):
	# arr은 할당된 상담요청
    standby_time = 0
    state_queue = []
    for a, b in arr:
    	# 상담가능 멘토가 없는 경우
        if m == 0:
        	# 가장 빨리 끝나는 상담 종료
            end_time = heapq.heappop(state_queue)
            m += 1
            # 상담이 늦게 끝난 경우
            if end_time > a:
            	# 대기시간 누적
                standby_time += end_time - a
                a = end_time
        m -= 1
        heapq.heappush(state_queue, a+b)
    return standby_time

상담가능 멘토가 없으면 상담 종료시각이 빠른 멘토부터 종료시키는 로직을 작성했다. 상담요청의 길이를 n이라 한다면 시간복잡도는 O(nlogn)이 된다.

이제 멘토를 할당할 방식을 생각해야 한다. 두 가지 방법을 생각해 볼 수 있다. 완전탐색과 그리디 알고리즘이다.

1) 완전탐색

멘토가 배치되는 모든 경우를 생각해보자. 최악의 경우를 상정하기 위해 유형은 5개, 멘토는 20명으로 설정한다. 각 유형별로 적어도 멘토가 1명씩 존재해야 하므로 15명의 멘토를 5개의 그룹으로 나누는 경우의 수를 구하면 된다. 이는 $19 \choose 4$ = 3876이 될 것이다.

또한, 유형별로 멘토의 수에 따른 최소 비용을 미리 캐싱해야 하므로 최악의 경우 1~20명의 멘토에 대해 300개의 상담요청에 대한 비용을 계산해야 한다. 이는 20 * 300log300로 계산 가능하며 대략 48000정도가 된다.

입력 데이터의 경우에는 크기가 작아서 완전탐색으로도 가능해 보이지만 데이터가 커질 경우 시간초과가 발생할 수 있다. 따라서 그리디 알고리즘을 선택하는 것이 더 좋은 선택인 것 같다.

2) 그리디 알고리즘

문제를 해결하는 것에 그리디 알고리즘이 적용되려면 다음 조건이 만족되어야 한다.

멘토가 n명에서 n+1명으로 증가할 때 감소하는 총 대기시간을 T(n)라 하면
T(n) >= T(n+1)을 항상 만족해야 한다.

위 조건을 만족한다면 현재 멘토의 배치에서 멘토 1명을 추가로 할당했을 경우의 최적이 전체 배치의 최적임이 자명하다. 멘토의 수와 대기시간은 반비례 관계에 있으므로 위 조건을 만족한다. 따라서 그리디 알고리즘을 적용할 수 있다.

import heapq
from collections import deque

def get_standby_time(arr, m):
	# arr은 할당된 상담요청
    standby_time = 0
    state_queue = []
    for a, b in arr:
    	# 상담가능 멘토가 없는 경우
        if m == 0:
        	# 가장 빨리 끝나는 상담 종료
            end_time = heapq.heappop(state_queue)
            m += 1
            # 상담이 늦게 끝난 경우
            if end_time > a:
            	# 대기시간 누적
                standby_time += end_time - a
                a = end_time
        m -= 1
        heapq.heappush(state_queue, a+b)
    return standby_time

def solution(k, n, reqs):
	# 각 유형별 상담요청
    counsel = [[] for _ in range(k)]
    
    for a, b, c in reqs:
        counsel[c-1].append([a, b])
    # index 0은 현재 대기시간, 1은 멘토 추가시 대기시간
    # 멘토가 1명일 경우가 기본값
    standby_time_arr = [[get_standby_time(counsel[i], 1) for i in range(k)], [0]*k]
    time_gap = []
    
    for i in range(k):
    	# 만약 대기시간이 0일 경우 고려할 필요 없음
        if standby_time_arr[0][i] == 0:
            continue
        # 멘토가 2명일 경우
        standby_time = get_standby_time(counsel[i], 2)
        # time_gap의 요소는 [-gap, idx, m]
        # 대기시간이 긴 유형이 위로
        heapq.heappush(time_gap, [-(standby_time_arr[0][i] - standby_time), i, 2])
        standby_time_arr[1][i] = standby_time
        
    for _ in range(n-k):
        if not time_gap:
            break
        # 현재 가장 대기시간을 줄일 수 있는 멘토의 배치
        _, i, m = heapq.heappop(time_gap)
        standby_time = get_standby_time(counsel[i], m+1)
        heapq.heappush(time_gap, [-(standby_time_arr[1][i] - standby_time), i, m+1])
        standby_time_arr[0][i] = standby_time_arr[1][i]
        standby_time_arr[1][i] = standby_time
        
    return sum(standby_time_arr[0])

결과는 다음과 같다.