[백준/Python]21944: 문제 추천 시스템 Version 2

문제

https://www.acmicpc.net/problem/21944
백준 문제집 - 0x16강 - 이진 검색 트리

 

알고리즘: 이진 탐색 트리, 트리를 사용한 집합과 맵

---

풀이

 

21939: 문제 추천 시스템 Version 1 과 비슷하게 풀고 무한 TLE♾️에 빠졌던 문제다.
좀 더 무식하게(?) 풀어야 AC를 받을 수 있다.

 

일단 문제의 조건을 정리해보면 다음과 같다.

 

입력 데이터 범위: 1 <= 알고리즘, 난이도 <= 100, 1 <= 문제번호 <= 100,000

 

recommend: 알고리즘 분류가 G인 문제

- 난이도 높고, 번호가 큰 문제

- 난이도 낮고, 번호가 작은 문제

 

recommend2: 알고리즘 상관 X

- 난이도 높고, 번호가 큰 문제

- 난이도 낮고, 번호가 작은 문제

 

recommend3: 난이도 L 이상/미만인 문제, 알고리즘 상관 X

- L 이상: 번호가 작은 문제

- L 미만: 번호가 큰 문제

 

👉 알고리즘별 최대/최소, 난이도별 최대/최소, 문제의 정보(번호, 난이도, 알고리즘)이 필요한 상태

 

문제의 정보와 난이도별 최대/최소는 리스트를, 알고리즘별 최대/최소는 딕셔너리를 사용한다.
알고리즘 같은 경우 G에 해당되는 문제들만 탐색하면 되지만, 난이도는 조건값이 명확히 정해져있지 않기 때문이다.
명령 조건에 알맞은 문제를 찾을때까지 난이도 높음→낮음 or 낮음→높음 순으로 탐색해야한다.

 

따라서 아래와 같이 초기화해준 뒤 진행한다.

난이도별 최대/최소 - 범위(1 ~ 100)를 가정한 2차원 리스트
알고리즘별 최대/최소 - 딕셔너리
문제의 정보 - 범위(1 ~ 100000)를 가정한 2차원 리스트

# 난이도별 힙
l_max = [[] for _ in range(101)]  # 최대힙
l_min = [[] for _ in range(101)]  # 최소힙

# 그룹별 힙
g_max = {}  # 최대힙 - 필요할 때만 생성
g_min = {}  # 최소힙 - 필요할 때만 생성

# 문제 정보를 O(1)에 접근하기 위한 배열
problems = [0] * 100001  # [난이도, 그룹] 저장

• problems에 문제의 정보가 저장되어있으므로 난이도별 힙에는 문제의 번호만 저장
• 알고리즘은 난이도, 문제의 번호를 모두 따져야하므로 (난이도, 문제 번호)로 저장

 

문제 데이터를 집어넣으면 아래와 같은 형태가 된다.

난이도별 최대힙 = [[-문제 번호1], [-문제 번호2]...]
난이도별 최소힙 = [[문제 번호1], [문제 번호2]...]

알고리즘별 최대힙 = {알고리즘1: [(-난이도, -문제 번호)]...}
알고리즘별 최소힙 = {알고리즘1: [(난이도, 문제 번호)]...}

문제의 정보[문제 번호] = [[난이도, 알고리즘]...]

 

add는 위 형태 그대로 진행, solved는 해당 문제의 값을 0으로 변경해준다. (힙에는 이전 정보가 그대로 남아있는 상태)

그리고 메인 포인트인 recommend ~ 3 는 조건분기를 명확히 작성해줘야한다.

 

recommend

- 힙[G]가 남아있을때까지 진행

   - 힙에서 pop한 값 ➡️ L, P

   - P의 값이 0이거나 저장되어있는 값과 [L, G]가 다를경우 ➡️ 이미 solved되거나 변경된 문제이므로 continue

   - 원하는 값을 찾았다면 출력, 다시 힙에 추가하고 break

 

recommed2

- 값을 못찾았을 경우를 판단하기 위해 플래그 설정 ➡️ found = False

- 최대/최소에 따라 범위 설정 후 난이도별로 하나씩 탐색 (1 ~ 100 or 100 ~ 1)

   - 해당 난이도 L의 힙이 남아있을때까지 탐색

      - 힙에서 pop한 값 ➡️ P

      - P의 값이 0이거나 저장되어있는 값과 L이 다를경우 ➡️ 이미 solved되거나 변경된 문제이므로 continue

      - 원하는 값을 찾았다면 출력, 다시 힙에 추가

      - 플래그를 True로 변경 후 break

- 모든 탐색이 끝난후에도 플래그가 False라면 -1 출력

 

recommend3

- 값을 못찾았을 경우를 판단하기 위해 플래그 설정 ➡️ found = False

- L부터 or L까지로 범위 설정 후 난이도별로 하나씩 탐색 (L ~ 100 or 1 ~ L-1)

   - 해당 난이도의 힙이 남아있을때까지 탐색

      - 힙에서 pop한 값 ➡️ P

      - P의 값이 0이거나 저장되어있는 값과 다를경우 ➡️ 이미 solved되거나 변경된 문제이므로 continue

      - 원하는 값을 찾았다면 출력, 다시 힙에 추가

      - 플래그를 True로 변경 후 break

- 모든 탐색이 끝난후에도 플래그가 False라면 -1 출력

 

전체 코드

# 메모리: 73604KB / 시간: 416ms
from sys import stdin
from heapq import heappush, heappop


input = stdin.readline

# 난이도별 힙
l_max = [[] for _ in range(101)]  # 최대힙
l_min = [[] for _ in range(101)]  # 최소힙

# 그룹별 힙
g_max = {}  # 최대힙 - 필요할 때만 생성
g_min = {}  # 최소힙 - 필요할 때만 생성

# 문제 정보를 O(1)에 접근하기 위한 배열
problems = [0] * 100001  # [난이도, 그룹] 저장

# 주어진 문제부터 처리
N = int(input())
for _ in range(N):
    P, L, G = map(int, input().split())
    problems[P] = [L, G]  # 문제 정보 저장

    # 난이도별 힙에 추가
    heappush(l_max[L], -P)
    heappush(l_min[L], P)

    # 그룹별 힙에 추가
    if G not in g_max:
        g_max[G] = []
        g_min[G] = []
    heappush(g_max[G], (-L, -P))
    heappush(g_min[G], (L, P))

# 명령대로 처리
M = int(input())
for _ in range(M):
    cmd, *data = input().rstrip().split()

    # 1. add
    if cmd == "add":
        P, L, G = map(int, data)
        problems[P] = [L, G]  # 문제 정보 저장

        heappush(l_max[L], -P)
        heappush(l_min[L], P)

        if G not in g_max:
            g_max[G] = []
            g_min[G] = []
        heappush(g_max[G], (-L, -P))
        heappush(g_min[G], (L, P))

    # 2. solved
    elif cmd == "solved":
        P = int(data[0])
        problems[P] = 0  # 문제 삭제 표시

    # 3-1. recommend
    elif cmd == "recommend":
        G, x = map(int, data)

        if x == 1:
            while g_max[G]:
                L, P = heappop(g_max[G])
                P = -P

                # 삭제됐거나 정보가 변경된 경우
                if not problems[P] or problems[P] != [-L, G]:
                    continue

                print(P)
                heappush(g_max[G], (L, -P))
                break
        else:
            while g_min[G]:
                L, P = heappop(g_min[G])

                if not problems[P] or problems[P] != [L, G]:
                    continue

                print(P)
                heappush(g_min[G], (L, P))
                break

    # 3-2. recommend2
    elif cmd == "recommend2":
        x = int(data[0])
        found = False

        if x == 1:
            for L in range(100, -1, -1):
                while l_max[L]:
                    P = -heappop(l_max[L])

                    if not problems[P] or problems[P][0] != L:
                        continue

                    print(P)
                    heappush(l_max[L], -P)
                    found = True
                    break

                if found:
                    break
        else:
            for L in range(101):
                while l_min[L]:
                    P = heappop(l_min[L])

                    if not problems[P] or problems[P][0] != L:
                        continue

                    print(P)
                    heappush(l_min[L], P)
                    found = True
                    break

                if found:
                    break

    # 3-3. recommend3
    else:
        x, L = map(int, data)
        found = False

        if x == 1:
            for i in range(L, 101):
                while l_min[i]:
                    P = heappop(l_min[i])

                    if not problems[P] or problems[P][0] != i:
                        continue

                    print(P)
                    heappush(l_min[i], P)
                    found = True
                    break

                if found:
                    break
        else:
            for i in range(L-1, -1, -1):
                while l_max[i]:
                    P = -heappop(l_max[i])

                    if not problems[P] or problems[P][0] != i:
                        continue

                    print(P)
                    heappush(l_max[i], -P)
                    found = True
                    break

                if found:
                    break

        if not found:
            print(-1)

데이터 관리 방식이 관건인 문제였다. 사실 이 방법(여러개의 힙 모음으로 관리)을 떠올리긴 했었는데... 진짜 되는줄은 몰랐다^^;;

결국 다른 분의 풀이를 참고해서 풀긴 했지만 속이 다 후련하다. 시간초과의 늪에 영원히 갇힐 줄 알았다