[BOJ 2075] N번째 큰 수 (Python)

풀이

처음에는 Max Heap 기반 우선순위 큐를 만들어 N x N 표의 모든 원소를 우선순위 큐에 넣어주었고, N-1번 pop한 후 우선순위 큐의 첫 번째 원소를 출력하는 방식으로 접근하였다. 하지만, 메모리 초과가 났다. --> 메모리 초과가 난 이유는 처음부터 우선순위 큐에 모든 원소를 넣어주었기에 발생했다고 생각한다. N의 최대 크기는 1,500으로 맨 처음 우선순위 큐에는 1,500^2 개의 원소가 삽입될 것이다.)

문제의 포인트는 상위 N개만 담는 Min Heap 기반 우선순위 큐를 활용하는 것이다. 문제의 로직을 간략하게 설명하려 한다. 로직을 본 후 코드블록을 보면 이해하는데 도움이 될 것이라고 생각한다.

N x N 표를 board[N - 1][N - 1]라고 표현하겠다.

1. 맨 처음 먼저 Min Heap 기반 우선순위 큐에 board[0][0, 1, 2, ... , N - 1]의 원소들, 즉 1행의 원소들을 모두 우선순위 큐에 삽입한다.

2. board[1][0, 1, 2, ... , N - 1]의 원소들, 즉 2행의 원소들을 우선순위 큐에 삽입하는데, 큐에 삽입할 원소가 큐의 첫 번째 원소(q[0])보다 클 때 큐에 삽입해준다. (가장 큰 상위 N개의 원소를 구해주어야 하기 때문에 이러한 조건이 필요하다.)
   --> 우선순위 큐에는 1행 ~ 2행까지의 모든 원소들 중에서 상위 N개의 원소들을 오름차순으로 가지고 있게 된다.

3. board[2][0, 1, 2, ... , N - 1]의 원소들, 즉 3행의 원소들을 우선순위 큐에 삽입하는데, 큐에 삽입할 원소가 큐의 첫 번째 원소(q[0])보다 클 때 큐에 삽입해준다. (가장 큰 상위 N개의 원소를 구해주어야 하기 때문에 이러한 조건이 필요하다.)
   --> 우선순위 큐에는 1행 ~ 3행까지의 모든 원소들 중에서 상위 N개의 원소들을 오름차순으로 가지고 있게 된다.

4. 마지막 행까지 위 과정을 반복한 후 큐의 첫 번째 원소를 반환한다.

이러한 로직이 가능한 이유는 문제의 '채워진 수에는 한 가지 특징이 있는데, 모든 수는 자신의 한 칸 위에 있는 수보다 크다는 것이다.' 라는 조건때문에 가능하다고 생각한다. (온전히 본인의 생각이다!)

코드

# Python3 메모리 초과 / PyPy3 통과!
import sys
import heapq

input = sys.stdin.readline
N = int(input())
board = [list(map(int, input().split())) for _ in range(N)]

q = []
# 우선순위 큐에 먼저 첫 번째 행의 원소들을 넣어준다.
for b in board[0]:
    heapq.heappush(q, b)

# 2번째 행부터 마지막 행까지 상위 N개의 수를 계속 갱신한다.
for i in range(1, N):
    for b in board[i]:
        if q[0] < b:
            heapq.heappush(q, b)
            heapq.heappop(q)

print(q[0])