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

문제

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

풀이

• N×N의 표에 수 N^2개 채워져 있다.
• 모든 수는 자신의 한 칸 위에 있는 수보다 크다.
• N번째 큰 수 출력

heapq를 사용하는 문제
처음 문제를 접했을 때, '모든 값을 다 삽입한 후에 heappop()을 N번 반복하면 되는 것이 아닌가?' 라는 생각이 들어 바로 구현해봤다.

import sys, heapq
input = sys.stdin.readline

n = int(input())
arr = [list(map(int,input().split())) for _ in range(n)]

heap = []
for i in range(n):
    for j in range(n):
        heapq.heappush(heap, -arr[i][j]) # 최대 힙

answer = 0
for k in range(n):
    answer = -heapq.heappop(heap)
print(answer)

그러나..

보기 좋게 실패 😞

그래서 더 간단하게 생각을 해봤습니다.
바로, heapq의 특성을 이용하는 것입니다.

우선 리스트를 바로 선언하지 않고 한 줄씩 입력을 받으면서 heapq에 바로 넣어주는 것입니다(‼️)

넣어줄 때, N의 길이를 초과하게 되면 거기서 heappop()을 해주는 것이죠.
그렇게 되면 heapq 안에 있는 가장 작은 값이 pop될 것이고, 그렇게 되면 상위 N개의 수들만 남게 됩니다.

for _ in range(n):
    row = list(map(int,input().split())) # 입력 받을 행
    for r in row: # 행 검사
        heapq.heappush(heap, r) # 우선 heapq에 push
        if len(heap) > n: # 길이를 초과할 경우
            heapq.heappop(heap) # pop

이러면 N의 길이를 가진 heapq를 구할 수 있게 됩니다.
여기서 가장 작은 값을 출력하게 되면, 그것이 N번째로 큰 수가 됩니다.

코드

import sys, heapq
input = sys.stdin.readline

n = int(input())
heap = []

for _ in range(n):
    row = list(map(int,input().split()))
    for r in row:
        heapq.heappush(heap, r)
        if len(heap) > n:
            heapq.heappop(heap)
print(heap[0])