[Python/파이썬] 백준 2075번 N번째 큰 수

2075번: N번째 큰 수

문제

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

입력

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

출력

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

 

코드

import sys, heapq

n = int(sys.stdin.readline())
heap = []

for _ in range(n):
    for i in list(map(int, sys.stdin.readline().split())):
        # 힙에 이미 n개의 원소가 있을 때 
        # 힙에서 가장 작은 원소와 비교하여 
        # 더 크면 가장 작은 원소 pop해버리고 본인이 들어감
        if len(heap) == n:  
            if heap[0] < i:
                heapq.heappop(heap)
                heapq.heappush(heap, i)
        else:   # 아직 힙에 n개의 원소보다 덜 찼을 때는 그냥 push
            heapq.heappush(heap, i)
            
print(heap[0])

최소 힙을 이용하여 구하였다.

힙의 원소가 N개보다 적다면 그냥 heappush로 push해주고 힙의 원소의 개수가 N개가 되었다면 힙에서 가장 작은 원소인 0번 인덱스의 원소와 입력값을 매번 비교하여 입력값이 힙에서 가장 작은 값보다 크다면 가장 작은 값을 pop하고 입력값을 push해준다. 이렇게 입력값을 모두 받고 나면 힙에 있는 원소 중 가장 작은 값이 이 표에서 N번째로 큰 수일 것이다. 따라서 힙의 0번 인덱스를 출력해주면 된다.