17612. 쇼핑몰 (Python)

17612번: 쇼핑몰

 

풀이

 고객이 계산대에 들어갈 때와 나갈 때, 두 경우에 대해서 최소 힙을 사용한다. 

 

 먼저 계산을 하러 들어갈 때는 최소 힙에 (시간, 계산대, 고객 번호)와 같은 형태로 넣어줬다. 만약 같은 시간에 계산대가 비는 경우 번호가 작은 계산대에 현재 손님이 배정된다.

 계산을 마치고 나가는 최소 힙에는 (시간, -계산대, 고객 번호)와 같은 형태로 값을 넣어줬다. 만약 같은 시간에 계산을 마쳤을 경우에는 출구에 가까운 높은 번호 계산대의 고객부터 빠져나가게 된다. 

 

 추가적으로 처음에 계산대의 수만큼 계산 시작 최소 힙에 (0, 계산대, -1)과 같은 형태로 for문을 통해서 값을 넣어줬는데 이렇게 할 경우 시간 초과가 발생했다. 처음에 for문을 통해서 모든 계산대를 최소 힙에 넣어주는 방식 대신 다른 방법을 사용해야 시간 초과 발생을 방지할 수 있다.

 

코드

import sys
import heapq
input = sys.stdin.readline

if __name__ == '__main__':
    N, K = map(int, input().split())
    before = []  # (시간, 계산대, 회원번호)
    after = []  # (시간, -계산대, 회원번호)

    counterIdx = 1  # 카운터 번호
    for _ in range(N):
        person, w = map(int, input().split())  # 회원번호, 물건 개수

        if counterIdx <= K:  # 아직 고객을 받지 않은 카운터가 남아있는 경우
            # 바로 before 힙에 값을 넣어준다.
            heapq.heappush(before, (w, counterIdx, person))
            counterIdx += 1
        else:  # 모든 카운터에 손님이 있을 경우
            time, counter, id = heapq.heappop(before)
            # 계산대에서 나오는 고객을 after 힙에 넣어주고
            heapq.heappush(after, (time, -counter, id))
            # 새로운 고객을 before 힙에 넣어준다.
            heapq.heappush(before, (time + w, counter, person))

    while before:  # 남은 계산대에 있는 손님들을 모두 after 힙으로 이동시킨다.
        time, counter, id = heapq.heappop(before)
        heapq.heappush(after, (time, -counter, id))

    res, idx = 0, 1
    while after:  # after 힙에 있는 손님들을 차례대로 꺼내어 계산한다.
        _, _, id = heapq.heappop(after)
        res += idx * id
        idx += 1

    print(res)