N번째 큰수(백준 2075번)

문제

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

N=int(input())
N_heap=[]

for _ in range(N):
    tmp_list=list(map(int,sys.stdin.readline().strip().split()))
    for tmp in tmp_list:
        if len(N_heap)<N:
            heapq.heappush(N_heap,tmp)
        else:
            if N_heap[0]<tmp:
                heapq.heappop(N_heap)
                heapq.heappush(N_heap,tmp)
print(N_heap[0])

• heap의 크기가 N을 유지하도록 한다.
  • heap으로 구조를 변경하면 가장 작은 값은 heap[0]이 맞지만, 그 외의 순서는 숫자의 크기와 상관없다.
• 가장 작은 값인 N_heap[0]보다 비교하는 값이 크면, heappop()을 한 후 다시 heappush()를 해준다.
• 이렇게 하면, 가장 큰 수 5개가 heap에 남는다. heapq는 MinHeap이므로, N_heap[0]이 N번째로 큰 수가 된다.

---

추가적인 문법

• nlargest, nsmallest

# heap_q에서 가장 큰 3개의 원소가 담긴 리스트
print(heapq.nlargest(n=3, iterable=heap_q))

# heap_q에서 가장 작은 3개의 원소가 담긴 리스트
print(heapq.nsmallest(n=3, iterable=heap_q))

# heap_q
# [0, 1, 3, 10, 4]

# heap_q의 nlargest
# [10, 4, 3]

# heap_q의 nsmallest
# [0, 1, 3]