[Python] BOJ 백준 21939 문제 추천 시스템 Version 1

문제

---

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

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

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

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

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

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

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

---

풀이

---

노가다로 해결한 문제. 단순히 분기를 나누고 min, max값을 구해서 해결할 수 있다고 생각했는데, 어림도 없었다.

이런 저런 방법들을 전부 사용해보다가, 같은 스터디원이 heap을 쓰다가 포기했다는 것을 듣고 heap으로 구현해봤다.

아래는 초기 실패 코드

import sys
from collections import Counter

si = sys.stdin.readline
def MIIS(): return map(int, si().split())


n = int(si())

prob = {k: [] for k in range(101)}
idx = {}
for _ in range(n):
    num, difficulty = MIIS()
    prob[difficulty].append(num)
    idx[num] = difficulty

_max = max(x for x in prob if prob[x])
_min = min(x for x in prob if prob[x])
m = int(si())
ret = []

for _ in range(m):
    line = si().split()
    if line[0] == 'add':
        num, diff = map(int, line[1:])
        if diff < _min:
            _min = diff
        elif diff > _max:
            _max = diff
        prob[diff].append(num)
        idx[num] = diff
    elif line[0] == 'solved':
        num = int(line[1])
        if num in prob[_min]:
            prob[idx[num]].remove(num)
            _min = min(x for x in prob if prob[x])
        elif num in prob[_max]:
            prob[idx[num]].remove(num)
            _max = max(x for x in prob if prob[x])
        else:
            prob[idx[num]].remove(num)
    elif line[0] == 'recommend':
        if line[1] == '-1':
            ret.append(sorted(prob[_min])[0])
        else:
            ret.append(sorted(prob[_max])[-1])

for a in ret:
    print(a)

삽질의 흔적...

아래는 성공한 코드이다. 

defaultdict를 이용하면서 해결 했는 지 체크하고, 명령어에 따라 min_heap과 max_heap을 적절히 섞어가며 사용했다.

애를 먹었던 부분은 add하는 부분에서 같은 문제가 해결되고 다른 난이도로 다시 들어올 수 있다는 사실을 간과했던 것이다.

이를 막기 위해 add에서도 해결된 문제들을 전부 빼주는 방식으로 진행했더니 해결되었다. 

from collections import defaultdict
import sys
import heapq

si = sys.stdin.readline
def MIIS(): return map(int, si().split())


n = int(si())

# recommend 용
min_heap = []
max_heap = []
ret = []

visited = defaultdict()

for _ in range(n):
    num, difficulty = MIIS()
    heapq.heappush(min_heap, (difficulty, num))
    heapq.heappush(max_heap, (-difficulty, -num))
    visited[num] = False
m = int(si())
for _ in range(m):
    line = si().split()
    if line[0] == 'add':
        num, difficulty = map(int, line[1:])
        while visited[(min_heap[0][1])]:
            heapq.heappop(min_heap)
        while visited[-(max_heap[0][1])]:
            heapq.heappop(max_heap)
        heapq.heappush(min_heap, (difficulty, num))
        heapq.heappush(max_heap, (-difficulty, -num))
        visited[num] = False

    elif line[0] == 'solved':
        visited[int(line[1])] = True

    else:
        if line[1] == '1':
            while visited[-(max_heap[0][1])]:
                heapq.heappop(max_heap)
            ret.append(-max_heap[0][1])
        else:
            while visited[min_heap[0][1]]:
                heapq.heappop(min_heap)
            ret.append(min_heap[0][1])


for ans in ret:
    print(ans)