[백준] 2075 파이썬(python) : N번째 큰 수 - heapq

https://www.acmicpc.net/problem/2075

2075번: N번째 큰 수

 

 

파이썬에서 제공하는 heapq와 입력값이 NxN으로 들어온다는 것을 활용하여 푸는 문제다. 파이썬 heapq는 최소 힙을 제공하기 때문에 한 줄을 입력받고 heapq에 push하면 최소 값이 heap[0]에 정렬된다. 두번째 줄부터 입력값과 heap[0]을 계속 비교하면서 heap[0]가 작으면 빼고 입력값을 넣는 과정을 반복하면 자연스럽게 5개의 가장 큰 숫자들만 남게 된다. (비교하고 삽입, 삭제하는 과정에서 자동으로 최소 힙으로 정렬되기 때문에) 메모리 제한이 12MB로 상당히 적기 때문에 heap 리스트 안에 들어가는 원소를 n개로 조정해야 한다.

 

import sys, heapq
n = int(sys.stdin.readline())
heap = []
for _ in range(n):
    number = list(map(int, sys.stdin.readline().split())) #1
    if not heap: #2
        for num in number:
            heapq.heappush(heap, num)
    else: #3
        for num in number:
            if heap[0] < num: #4 
                heapq.heappush(heap, num)
                heapq.heappop(heap)
print(heap[0]) #5

 

#1 : 한 줄씩 입력받음

#2 : 비어있는 경우(첫 줄을 입력 받은 경우) 모두 입력 받고 최소 힙으로 정렬

#3, 4 : heap안에서 가장 작은 원소인 heap[0]을 들어오는 원소와 크기 비교해서 작다면 빼고 들어오는 원소 삽입 heap[0] < num이어야 삽입, 삭제가 이뤄지기 때문에 가장 큰 숫자들만 남게 됨

#5 : n개의 가장 큰 숫자들 중 가장 작은 숫자가 n * n개의 숫자들 중 n번째로 큰 숫자이므로 heap[0]를 출력