N×N의 표에 수 N^2개 채워져 있다. 채워진 수에는 한 가지 특징이 있는데, 모든 수는 자신의 한 칸 위에 있는 수보다 크다는 것이다. 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번째 큰 수를 찾는 프로그램을 작성하시오. 표에 채워진 수는 모두 다르다.
이번 문제의 핵심은 메모리 초과이다.
즉, 기존에 배열을 사용해서 정렬을 진행하면 메모리 초과가 난다.
배열이 아닌 힙 형태로 문제를 해결할 수 있다.
1. 첫번재 행으로 최소 힙을 채운다.
2. 2번째 행부터 각 숫자가 최소 힙의 루트보다 크다면 최소 힙에 집어넣는다.
3. 최소 힙에 들어간 숫자는 루트보다 크기 때문에 마지막 레벨에 리프 노드로 붙여진다.
4. 루트 노드를 pop하면 그 다음으로 작은 노드가 루트 노드가 되고 힙은 다시 N의 노드를 갖는다.
5. 2 ~ 4 과정을 반복한다.
import heapq
import sys
if __name__ == '__main__':
N = int(input())
min_heap = []
for _ in range(N):
row = map(int, sys.stdin.readline().rstrip().split())
if not min_heap:
for num in row:
heapq.heappush(min_heap, num)
else:
for num in row:
if min_heap[0] < num:
heapq.heappush(min_heap, num)
heapq.heappop(min_heap)
print(min_heap[0])
원래 처음에는 우선순위 큐를 활용해 문제를 풀 예정이었다.
하지만 파이썬의 우선순위 큐 모듈을 쓰니 시간 초과가 발생해서 최소 힙으로 문제를 풀었다.
우선순위 큐도 바이너리 트리로 구현되어 있을 것 같은데 시간 초과가 발생한 이유는 잘 모르겠다... heapq보다 더 복잡한 내부 구현이 되어있는 듯 싶다.
사실 heapq 자체도 832ms로 아슬아슬하게 통과했다. 이럴때는 c++로 코드를 짜야되나 많이 고민되는 문제였다 ㅠㅠ
아무튼 골드 문제여서 많이 쫄면서 풀었는데 메모리와 시간 초과가 발생하는 문제만 해결하면 비교적 간단한 문제였다.