[알고리즘] 입국심사에서 각 입국 심사자의 평균 대기 시간 구하기(심사대 k개, 최소힙 이용)

문제 기술

프로그래머스 lv3 입국심사 문제를 변형한 문제이다.

공항에서 사람들의 입국을 심사하는 심사대가 k개있다. 사람들이 입국심사를 받기 위하여, 심사대 바로 앞에서 한 줄로 서서 기다린다. k개의 심사대 중 하나가 비어있으면(심사받는 사람이 없으면), 먼저 도착한 사람이 비어 있는 이 심사대에서 심사를 받는다. 입국심사를 받기 위한 각 사람의 심사대 앞에 도착하는 시간과 심사받는데 걸리는 시간이 주어진다. 각 사람이 심사대 앞에 도착한 시간에 대한 정보는, 바로 이전 사람이 심사대 앞에 도착한 시간과의 차이(분)로 주어진다. 각 사람이 심사를 받기 위하여 기다리는 시간 (심사시간은 제외)의 평균을 소수점 이하 1자리까지 (소수점 이하 2째 자리에서 반올림) 출력하는 프로그램을 작성하시오.

예를 들어, k = 1이라 하자. 그리고 5명의 사람들의 도착시간이 0, 2, 1, 2, 1이고, 심사받는데 걸리는 시간이 각각 4분, 7분, 6분, 5분, 3분이라 하자. 5명의 도착시간이 0, 2, 1, 2, 1라는 것은, 2번째 사람은 첫 번째 사람이 심사대 앞에 도착한 시간 2분 후에 심사대 앞에 도착하고, 3번째 사람은 2번째 사람이 도착한 시간 1분 후에 심사대 앞에 도착하고, 4번째 사람은 3번째 사람이 도착한 시간 2분 후에 심사대 앞에 도착하고, 5번째 사람은 4번째 사람이 도착한 시간 1분 후에 심사대 앞에 도착한다.

첫 번째 사람은 기다리지 않고 심사를 받는다. 두 번째 사람은 첫 번째 사람이 심사를 받은 직후, 심사를 받으므로 2분을 기다린다. 세 번째, 4번째, 5번째 사람은 각각 8분, 12분, 16분 기다린다. 따라서 5명이 심사대 앞에서 기다리는 평균시간은 (2+8+12+16)/5 =38/5 = 7.6분이다.

위의 예에서 k = 2라 하자. 심사대가 2개이므로 첫 번째 사람과 두 번째 사람은 각각 심사대 1과 심사대 2에서 기다리지 않고 심사를 받는다. 3번째 사람은 1분을 기다리고 심사대 1에서 심사를 받는다. 4번째 사람은 4분을 기다리고 심사대 2에서 심사를 받는다. 5번째 사람은 4분을 기다리고 심사대 1에서 심사를 받는다. 따라서 5명이 심사대 앞에서 기다리는 평균시간은 (0 + 0 + 1 + 4 + 4) / 5 = 1.8분이다.

요구조건: k개 심사대를 최소힙을 이용하여야 함

입력

첫 번째 줄에 심사대 수 k(1 이상 1000 이하 양의 정수)가 주어진다. 두 번째 줄에 심사대 앞에 도착하는 사람들의 수 n이 주어진다. 다음의 각 줄에 먼저 온 사람부터 각 사람의 심사대 앞 도착시간(분 단위 양의 정수)과 심사받는데 걸리는 시간(분 단위의 양의 정수)이 주어진다.

출력

각 사람의 평균 대기 시간이 출력된다.

입력 예 1

1
5
0 4
2 7
1 6
2 5
1 3

출력 예 1

7.6

입력 예 2

1
3
0 3
2 4
3 3

출력 예 2

1.0

입력 예 3

2
5
0 4
2 7
1 6
2 5
1 3

출력 예 3

1.8

입력 예 4

2
3
0 3

출력 예 4

0.0

나의 코드

import heapq
import sys

arvt = []
judgt = []
min_heap = []
n = int(sys.stdin.readline())
k = int(sys.stdin.readline())

for i in range(k):
   arvt_idx, judgt_idx = map(int, sys.stdin.readline().split())
   arvt.insert(i, arvt_idx)
   judgt.insert(i, judgt_idx)  # 여기까지 입력값을 리스트로 만드는 과정.
   # ex) 5 / 0 4 / 2 7 / 1 6 / 2 5 / 1 3일 경우
   # arvt = [0, 2, 1, 2, 1], judgt = [4, 7, 6, 5, 3]

total_time_at_immigration = [0 for i in range(k + 1)]  # 심사대에서 총 걸린 시간
pure_wait_time = [0 for i in range(k+1)]  # 우리가 구해야 할 값(각 심사자의 대기시간)
sum_arvt = [0 for i in range(k + 1)]  # 각 심사자의 누적 도착 시간

num = 0
for i in range(len(arvt)):
   if arvt[i] != 0:  # 도착 시간이 0이 아닐 경우에
       num += arvt[i]  # 그 값을 num에 추가하고,
       sum_arvt[i] += num  # 그 num을 sum_arvt[i]에 추가한다.
       # 이 if문이 끝나면 sum_arvt에는 각 참가자별 누적 도착 시간이 들어가게 된다.
for i in range(n):  # 심사대가 n개니까, n명의 사람들은
   heapq.heappush(min_heap, sum_arvt[i] + judgt[i])
   # 우선 대기자가 생기기 전까지는 min_heap에 누적 도착 시간+심사시간 만 넣는다.

for i in range(n, k):
   if sum_arvt[i] >= min_heap[0]:
       # min_heap은 최소힙 처리가 되어있으므로, root 노드인 0번째 인덱스가 제일 작은 값.
       # 그 다음 온 사람의 누적 도착 시간이 그 전 누군가의 누적 도착 시간 + 심사시간(제일 짧음)보다 클 경우,
       # 즉 심사대가 1개 이상 비어있을 경우이다.
       # 그 사람은 대기 시간이 없다.
       total_time_at_immigration[i] = sum_arvt[i] + judgt[i]
       pure_wait_time[i] = 0  # 그래서, 이 심사자의 대기시간은 0분이다.
   else:  # 심사대가 비어있지 않을 경우에는,
       total_time_at_immigration[i] = sum_arvt[i] + abs(sum_arvt[i] - min_heap[0]) + judgt[i]
       # i번째 사람의 대기시간은 (누적 도착 시간 + 심사시간) + 절댓값(도착시간 - 가장 짧은 심사자의 누적도착시간 + 심사시간)
       pure_wait_time[i] = abs(sum_arvt[i] - min_heap[0])
       # 순수 대기 시간은, 본인의 누적 도착 시간에서 가장 짧은 심사자가 총 머문 시간을 뺀 값에 절댓값을 붙인 값이다.
   heapq.heappop(min_heap)  # 이제 가장 짧은 사람은 심사를 마쳤으므로 heapq.heappop을 통해 제거한다.
   # 그리고, i번째 사람이 심사대에 들어가게 되었으므로 heapq.heappush를 통해 추가한다.
   heapq.heappush(min_heap, total_time_at_immigration[i])

print(round(sum(pure_wait_time)/k, 1))  # 순수 대기 시간을 전부 더한 뒤, 소수점 둘째 자리에서 반올림한다.

문제를 풀어나가면서 알고리즘도 재밌을 수가 있구나 하고 느꼈다. 이 문제는 풀어가는 과정이 되게 재밌었다. 힙이 어떻게 형성되고 바뀌는지를 머릿속에서 생각해 나가면서 풀었다. 몇시간 걸렸던 것으로 기억하지만, 재밌었다.
x번째 사람이 왔을 때 심사대가 비어있는지 비어있지 않은지를 기준으로 크게 두 가지로 나눴다. 심사대가 비어있다면 그냥 들어가면 되니까 대기시간이 발생하지 않고, 심사대가 비어있지 않다면 x번째 사람의 대기시간은 절댓값(x번째 사람의 도착 시간 - 가장 짧은 심사자의 누적 도착 시간)이 된다.