문제

N×N의 표에 수 N2개 채워져 있다. 채워진 수에는 한 가지 특징이 있는데, 모든 수는 자신의 한 칸 위에 있는 수보다 크다는 것이다. N=5일 때의 예를 보자.
예

입력
12	7	9	15	5
13	8	11	19	6
21	10	26	31	16
48	14	28	35	25
52	20	32	41	49

이러한 표가 주어졌을 때, N번째 큰 수를 찾는 프로그램을 작성하시오.

코드

import heapq, sys

input = sys.stdin.readline
n = int(input())
heap = list()

for i in range(n):
    arr = list(map(int, input().split()))
    for j in arr:
        heapq.heappush(heap, j)
        if len(heap) > n:
            heapq.heappop(heap)

print(heap[0])

설명

처음에는 입력받은 n^2개의 수를 모두 힙에 저장해서 n번 pop했을 때의 수를 출력했는데 메모리 초과가 났다. 그래서 고민하다가 결국 검색을 해봤는데 최소 힙은 알다시피 힙에 5개의 숫자가 저장되어 있다면 최상단의 숫자는 그 n번째로 큰 수이다. 즉 입력의 처음부터 끝까지 숫자를 하나씩 저장하면서 힙의 크기가 n을 넘어가면 pop하고 아니면 push해준다면 입력이 마친 후엔 힙의 최상단에 남는 숫자는 n번째로 큰 수가 된다.