[Python] 백준 21939 - 문제 추천 시스템 Version 1 문제 풀이

Overview

BOJ 21939번 문제 추천 시스템 Version 1 Python 문제 풀이
분류: Data Structure (자료구조)

---

문제 페이지

https://www.acmicpc.net/problem/21939

---

풀이 코드 1

from sys import stdin
from collections import defaultdict
import heapq


class Problems:
    def __init__(self):
        self.in_list = defaultdict(bool)
        self.easy = []
        self.hard = []

    def recommend(self, x) -> int:
        if x == 1:
            while not self.in_list[-self.hard[0][1]]:
                heapq.heappop(self.hard)
            return -self.hard[0][1]
        else:
            while not self.in_list[self.easy[0][1]]:
                heapq.heappop(self.easy)
            return self.easy[0][1]

    def add(self, p, l) -> None:
        self.in_list[p] = True
        heapq.heappush(self.easy, (l, p))
        heapq.heappush(self.hard, (-l, -p))

    def solved(self, p) -> None:
        self.in_list[p] = False
        while self.easy and not self.in_list[self.easy[0][1]]:
            heapq.heappop(self.easy)
        while self.hard and not self.in_list[-self.hard[0][1]]:
            heapq.heappop(self.hard)


def main():
    def input():
        return stdin.readline().rstrip()

    prob = Problems()
    n = int(input())
    for _ in range(n):
        p, l = map(int, input().split())
        prob.add(p, l)
    m = int(input())
    for _ in range(m):
        cmd = list(input().split())
        if cmd[0] == 'add':
            prob.add(int(cmd[1]), int(cmd[2]))
        elif cmd[0] == 'recommend':
            print(prob.recommend(int(cmd[1])))
        else:
            prob.solved(int(cmd[1]))


if __name__ == "__main__":
    main()

난이도 쉬운 순서로 저장한 heapq easy와 난이도 어려운 순서로 저장한 heapq hard 에 문제들을 저장하여 구현한다. 그리고 in_list defaultdict에 각 문제가 문제 리스트에 들어있는지 여부를 저장한다.

---

풀이 코드 2 - python3 시간초과, pypy3 통과

from sys import stdin
from collections import defaultdict


class Problems:
    def __init__(self):
        self.levels = defaultdict(list)
        self.problems = {}

    def recommend(self, x) -> int:
        if x == 1:
            return max(self.levels[max(self.levels.keys())])
        else:
            return min(self.levels[min(self.levels.keys())])

    def add(self, p, l) -> None:
        if p in self.problems:
            self.levels[self.problems[p]].remove(p)
        self.problems[p] = l
        self.levels[l].append(p)

    def solved(self, p) -> None:
        self.levels[self.problems[p]].remove(p)
        if not self.levels[self.problems[p]]:
            self.levels.pop(self.problems[p])
        self.problems.pop(p)


def main():
    def input():
        return stdin.readline().rstrip()

    prob = Problems()
    n = int(input())
    for _ in range(n):
        p, l = map(int, input().split())
        prob.add(p, l)
    m = int(input())
    for _ in range(m):
        cmd = list(input().split())
        if cmd[0] == 'add':
            prob.add(int(cmd[1]), int(cmd[2]))
        elif cmd[0] == 'recommend':
            print(prob.recommend(int(cmd[1])))
        else:
            prob.solved(int(cmd[1]))


if __name__ == "__main__":
    main()
    

pypy3 채점은 통과하지만, python3 채점은 시간초과가 발생하는 코드이다. 힙을 사용하지 않고 defaultdict에 문제들을 저장하여 구현하였다.