문제

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번째 큰 수를 찾는 프로그램을 작성하시오. 표에 채워진 수는 모두 다르다.

www.acmicpc.net/problem/2075

  • 입력

첫째 줄에 N(1 ≤ N ≤ 1,500)이 주어진다. 다음 N개의 줄에는 각 줄마다 N개의 수가 주어진다. 표에 적힌 수는 -10억보다 크거나 같고, 10억보다 작거나 같은 정수이다.

  • 출력

첫째 줄에 N번째 큰 수를 출력한다.

최종코드

얼핏보면 쉬워보이지만 제한된 메모리 내에서 해결할 방법을 찾아야 했다. N번째 큰 수를 찾는 것이 목적이므로 min heap을 사용하고 min heap의 크기를 최대 2*N으로 제한해서 풀 수 있다.

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

위와 같은 예시를 살펴보자. 우선 맨 첫 번째 줄은 전부 min heap에 넣는다. 편의상 오름차순 리스트로 표현했다.

  • line 1 push [5, 7, 9, 12, 15]

두 번째 줄 부터는 한 줄을 min heap에 추가하고 N만큼 pop하는 것을 반복한다.

  • line 2 push [5, 6, 7, 8, 9, 11, 12, 13, 15, 19]
  • n개 pop [11, 12, 13, 15, 19]
  • line 3 push [10, 11, 12, 13, 15, 16, 19, 21, 26, 31]
  • n개 pop [16, 19, 21, 26, 31]
  • line 4 push [14, 16, 19, 21, 25, 26, 28, 31, 35, 48]
  • n개 pop [26, 28, 31, 35, 48]
  • line 5 push [20, 26, 28, 31, 32, 35, 41, 48, 49, 52]
  • n개 pop [35, 41, 48, 49, 52]

이렇게 위 과정을 반복하면 heap에는 가장 큰 수 n개만 남게 된다. 여기서 pop을 한 번 더 하면 n번째로 큰 수를 구할 수 있다.

import sys
import heapq

n = int(sys.stdin.readline())
heap = []
for i in range(n):
    arr = list(map(int, sys.stdin.readline().split()))
    for j in range(n):
        heapq.heappush(heap, arr[j])
    if i != 0:
        for j in range(n):
            heapq.heappop(heap)

print(heapq.heappop(heap))

Enhanced

위에서는 한 줄씩 push와 pop을 반복해서 2*N 사이즈의 heap이 필요했지만 push와 pop을 한 번씩 반복하면 N+1 사이즈의 heap으로도 문제를 해결할 수 있다.

import sys
import heapq

n = int(sys.stdin.readline())
heap = []
for i in range(n):
    arr = list(map(int, sys.stdin.readline().split()))
    if i == 0:
        for j in range(n):
            heapq.heappush(heap, arr[j])
    else:
        for j in range(n):
            # arr[j]가 heap[0]보다 작은 경우 어차피 바로 pop해야 하므로 처리하지 않고 넘긴다
            if heap[0] < arr[j]:
                heapq.heappush(heap, arr[j])
                heapq.heappop(heap)

print(heapq.heappop(heap))

push와 pop을 한 번에 해주는 heappushpop(heap, number) 함수를 사용할 수도 있다.

그리고 위의 풀이들은 문제의 조건을 전혀 활용하지 않았는데, 조건을 활용한 풀이가 궁금하다면 두 번쨰 참고 사이트를 확인하면 된다.

참고 사이트