문제 추천 시스템 Version 1(백준 21939번)

문제

tony9402는 최근 깃헙에 코딩테스트 대비 문제를 직접 뽑아서 “문제 번호, 난이도“로 정리해놨다.

깃헙을 이용하여 공부하시는 분들을 위해 새로운 기능을 추가해보려고 한다.

만들려고 하는 명령어는 총 3가지가 있다. 아래 표는 각 명령어에 대한 설명이다.

| recommend x | x가 1인 경우 추천 문제 리스트에서 가장 어려운 문제의 번호를 출력한다. 만약 가장 어려운 문제가 여러 개라면 문제 번호가 큰 것으로 출력한다.  x가 -1인 경우 추천 문제 리스트에서 가장 쉬운 문제의 번호를 출력한다. 만약 가장 쉬운 문제가 여러 개라면 문제 번호가 작은 것으로 출력한다. | | — | — | | add P L | 추천 문제 리스트에 난이도가 L인 문제 번호 P를 추가한다. (추천 문제 리스트에 없는 문제 번호 P만 입력으로 주어진다. 이전에 추천 문제 리스트에 있던 문제 번호가 다른 난이도로 다시 들어 올 수 있다.) | | solved P | 추천 문제 리스트에서 문제 번호 P를 제거한다. (추천 문제 리스트에 있는 문제 번호 P만 입력으로 주어진다.) |

명령어 recommend는 추천 문제 리스트에 문제가 하나 이상 있을 때만 주어진다.

명령어 solved는 추천 문제 리스트에 문제 번호가 하나 이상 있을 때만 주어진다.

위 명령어들을 수행하는 추천 시스템을 만들어보자.

나의 풀이

from heapq import heappop,heappush
import sys

N=int(input())
P_book=[False]*100001
Min_Heap=[]
Max_Heap=[]
for _ in range(N):
    I=sys.stdin.readline().strip().split()
    P_book[int(I[0])]=True
    heappush(Min_Heap,(int(I[1]),int(I[0])))
    heappush(Max_Heap,(-int(I[1]),-int(I[0])))
M=int(input())
count=N
for _ in range(M):
    cmd=sys.stdin.readline().strip().split()
    if cmd[0]=='recommend':
        tmp_list=[]
        if cmd[1]=='1' and count>0:
            while not P_book[-Max_Heap[0][1]]:
                heappop(Max_Heap)
            print(-Max_Heap[0][1])
        elif cmd[1]=='-1' and count>0:
            while not P_book[Min_Heap[0][1]]:
                heappop(Min_Heap)
            print(Min_Heap[0][1])
    elif cmd[0]=='add':
        while not P_book[-Max_Heap[0][1]]:
            heappop(Max_Heap)
        while not P_book[Min_Heap[0][1]]:
            heappop(Min_Heap)
        P_book[int(cmd[1])]=True
        heappush(Min_Heap,(int(cmd[2]),int(cmd[1])))
        heappush(Max_Heap,(-int(cmd[2]),-int(cmd[1])))
        count+=1
    elif cmd[0]=='solved':
        P_book[int(cmd[1])]=False
        count-=1
        if count==0:
            Min_Heap=[]
            Max_Heap=[]
                

• 이번 문제는 2개의 MinHeap, MaxHeap을 사용하여 푸는 것까지는 잘 접근했다. 하지만 한가지 몰랐던 것 때문에 정답을 찾는 것이 길어졌다.
  • 튜플로 구성된 리스트를 sort()했을때, 튜플의 첫번째 원소가 같으면 두번째 원소를 비교하여 정렬하는 것처럼, heapq도 마찬가지로 정렬된다.
  • 따라서 원래 같은 난이도를 갖는 문제를 다 뽑아서, 가장 큰 값을 출력한 후 heap에 다시 넣어주는 코드를 짰는데, 그럴 필요가 없다.