99클럽 코테 스터디 20일차 TIL + 우선순위 큐, defaultdict()

1. 문제 설명

[문제 내용]

회전 초밥 가게에 $N$명의 손님이 있고, 요리사는 $M$개의 초밥을 순서대로 만든다. 요리사가 초밥을 만들 경우, $1$번 손님부터 $N$번 손님의 순서대로 그 초밥을 받게 된다. 만약 먼저 초밥을 받는 손님이 초밥을 먹을 경우, 뒤의 손님들은 해당 초밥을 먹을 수 없다. 만약 아무도 해당 초밥을 먹지 않는다면, 초밥은 버려진다.

$N$명의 손님은 각자 먹고 싶은 초밥이 적힌 주문 목록을 가지고 있다. 목록에 적힌 초밥의 순서에 상관 없이 만약 목록에 적혀있는 초밥이 앞에 오면 반드시 먹는다. 만약, 목록에 적히지 않은 초밥을 받는다면 그 초밥은 반드시 먹지 않는다. 단, 손님들은 다양한 초밥을 먹고 싶어하기 때문에 각 종류의 초밥은 최대 한 번만 먹는다.

각 손님의 주문 목록과 순서대로 만들어지는 $M$개의 초밥이 주어질 때, 각 손님이 먹게 되는 초밥의 개수를 구해 보자.

[입력]

첫 번째 줄에 손님의 수 $N$과 초밥의 수 $M$이 공백으로 구분되어 주어진다. ($1\le N\le 100\ 000; 1\leq M\leq 200\, 000$)

두 번째 줄부터 $N$개의 줄에 걸쳐 각 손님에 대한 주문 목록을 나타내는 정수 $k$와 $A_1,A_2,\ldots, A_k$가 공백으로 구분되어 순서대로 주어진다. $k$는 주문 목록에 적힌 초밥 종류의 개수이며, $A_i$는 주문 목록에 적힌 초밥 종류이다. ($1\leq A_i\leq 200\, 000$)

각 손님의 주문 목록에 적힌 초밥 종류는 모두 다르며, 모든 손님에 대해 $k$의 합이 $200\, 000$ 이하임이 보장된다.

$N+2$번째 줄에 요리되는 초밥의 종류를 나타내는 $M$개의 정수 $B_1,B_2,\ldots, B_M$이 공백으로 구분되어 순서대로 주어진다. ($1\leq B_i\leq 200\, 000$)

[출력]

한 줄에 $N$개의 정수 $C_1,C_2,\ldots C_N$을 공백으로 구분하여 출력한다. $C_i$는 $i$번 손님이 먹은 초밥의 개수를 의미한다.

[입출력 예]

먼저 2번 손님이 3번 초밥과 2번 초밥을 먹는다. 이후, 1번 초밥이 제공되는데 1번 손님이 먼저 이를 먹기 때문에 3번 손님은 본인의 선호 목록에 1번 초밥이 있어도 이를 먹지 못한다. 이후 4번 초밥이 들어오는데 선호 목록에 4번 초밥이 있는 손님이 없으므로 바로 버려진다. 마지막으로 2번 손님이 5번 초밥을 먹게 된다.

따라서 1번 손님은 1개, 2번 손님은 3개, 3번 손님은 0개의 초밥을 먹게 된다.

2. 풀이

import sys
from collections import defaultdict

N, M = map(int, sys.stdin.readline().split())

# 초밥을 원하는 손님 목록을 저장할 딕셔너리
sushi_preferences = defaultdict(list)
eaten_sushi = [set() for _ in range(N)] # 각 손님이 먹은 초밥 저장
sushi_count = [0] * N # 각 손님이 먹은 초밥 개수

for i in range(N):
    data = list(map(int, sys.stdin.readline().split()))
    k, favorites = data[0], data[1:]
    
    for sushi in favorites:
        sushi_preferences[sushi].append(i) # 해당 초밥을 원하는 손님 리스트에 추가

sushi_list = list(map(int, sys.stdin.readline().split()))

# 초밥을 순서대로 처리
for sushi in sushi_list:

    if sushi in sushi_preferences: # 이 초밥을 원하는 손님이 있을 때만 확인
    
        for customer in sushi_preferences[sushi]: # 초밥을 원하는 손님 목록 불러오기
        
            if sushi not in eaten_sushi[customer]: # 이미 먹지 않았다면
                sushi_count[customer] += 1 # 먹은 초밥 개수 증가
                eaten_sushi[customer].add(sushi) # 초밥 먹었다고 기록
                break # 가장 앞 순서의 손님이 먹었으므로 종료

print(" ". join(map(str, sushi_count)))

3. 회고

3-1. 문제 해결 과정

이번 문제는 시간 초과가 발생해서 당황했다. 초밥이 나올 때마다 모든 손님이 먹을 것인지 확인하니 시간이 오래 걸렸던 거 같다. 그렇다면 초밥이 나올 때 이 초밥을 원하는 손님을 바로 찾을 수는 없을까? 고민을 하게 되었다.

# 초밥을 원하는 손님 목록을 저장할 딕셔너리
sushi_preferences = defaultdict(list)
eaten_sushi = [set() for _ in range(N)] # 각 손님이 먹은 초밥 저장
sushi_count = [0] * N # 각 손님이 먹은 초밥 개수

이전 시간에 배운 defaultdict()를 이용해서 초밥 종류를 기준으로 해당 초밥을 원하는 손님을 기록한다. eaten_sushi에는 각 손님이 먹은 초밥을 저장하고, sushi_count에는 각 손님이 먹은 초밥 개수를 기록한다.

# N명의 손님이 각자 먹고 싶은 초밥 목록을 입력 받음
for i in range(N):
    data = list(map(int, sys.stdin.readline().split()))
    k, favorites = data[0], data[1:]
    
    for sushi in favorites:
        sushi_preferences[sushi].append(i) # 해당 초밥을 원하는 손님 리스트에 추가

그 다음에는 k에 손님이 원하는 초밥의 개수를, favorites에는 손님이 원하는 초밥의 종류를 저장한다. 그 뒤에는 sushi_preferences에 해당 초밥을 원하는 손님을 기록했다.

# 초밥을 순서대로 처리
for sushi in sushi_list:

    if sushi in sushi_preferences: # 이 초밥을 원하는 손님이 있을 때만 확인

        for customer in sushi_preferences[sushi]: # 초밥을 원하는 손님 목록 불러오기

            if sushi not in eaten_sushi[customer]: # 이미 먹지 않았다면
                sushi_count[customer] += 1 # 먹은 초밥 개수 증가
                eaten_sushi[customer].add(sushi) # 초밥 먹었다고 기록
                break # 가장 앞 순서의 손님이 먹었으므로 종료

요리된 초밥을 입력 받은 다음에는 sushi_preferences에 해당 초밥이 있는지 없는지 확인하면서 진행한다. 해당 초밥을 원하는 손님 목록을 불러와서, 손님이 먹으면 해당 손님의 sushi_count를 증가시킨다. 그리고 손님이 먹은 초밥을 eaten_sushi에 기록한다.

3-2. 새롭게 배운 내용

• sushi_count = [0] * N : [0]이라는 하나짜리 리스트를 만들고 이걸 N번 반복해서 새로운 리스트를 만들겠다는 의미이다. 이 문제에서는 N명의 손님이 초밥을 몇개 먹었는지 확인하기 위해 사용되었다.

N = 3
sushi_count = [0] * N  # [0, 0, 0] (초기 상태)

sushi_count[1] += 1  # 2번 손님이 초밥을 하나 먹음
print(sushi_count)  # [0, 1, 0]

sushi_count[1] += 1  # 2번 손님이 또 하나 먹음
print(sushi_count)  # [0, 2, 0]

sushi_count[0] += 1  # 1번 손님이 초밥을 하나 먹음
print(sushi_count)  # [1, 2, 0]