[백준] #2075 N번째 큰 수 (python)

2075번: N번째 큰 수

2075번: N번째 큰 수

문제

N×N의 표에 수 N2개 채워져 있다. 채워진 수에는 한 가지 특징이 있는데, 모든 수는 자신의 한 칸 위에 있는 수보다 크다는 것이다. N=5일 때의 예를 보자.

이러한 표가 주어졌을 때, N번째 큰 수를 찾는 프로그램을 작성하시오. 표에 채워진 수는 모두 다르다.

입력

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

출력

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

예제 입력 1

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

예제 출력 1

35

❌ Code 1 (메모리 초과)

import heapq
import sys
input = sys.stdin.readline

n = int(input())
nums = [list(map(int, input().split())) for _ in range(n)]
heap = []
for i in nums:
    for j in i:
        heapq.heappush(heap, (-j, j))

for i in range(n-1):
    heapq.heappop(heap)

print(heap[0][1])

 

처음에는 nums 라는 리스트에 n*n 개의 수를 모두 집어넣고 max heap 에 넣어서 n-1번 pop 하면서 값을 구하는 식으로 구현을 했었다. 그런데 이렇게 하니… 메모리 초과가 뜬다. 검색해보니 힙의 갯수를 일정하게 유지하면서 풀어야 메모리 초과가 안 난다고 한다.

✔️ Code 2 (참고)

import heapq
import sys
input = sys.stdin.readline

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

for _ in range(n):
    nums = list(map(int,sys.stdin.readline().split()))

    if not heap:
        for num in nums: 
            heapq.heappush(heap, num)
        
    else:
        for num in nums:
            if num > heap[0]:
                heapq.heappop(heap)
                heapq.heappush(heap, num)
              
               
print(heap[0])

 

nums를 표의 모든 값을 저장하는 리스트로 만들어두지 않고, 한 줄을 입력받는 변수로 바꿨다.

그리고 max heap이 아닌 min heap을 사용했다.

N번째 큰 숫자를 구해야 하기 때문에 힙의 크기를 N으로 유지하는 것이 중요하다.

대충 로직은 이렇다.

5번째로 큰 수를 구해야 한다면, 우선 입력받은 맨 처음 5개 수를 모두 heap에 집어넣는다. (초기 heap은 비어있는 상태이기 때문에)

min heap이기 때문에 자동으로 가장 작은 수가 heap의 첫 번째 원소로 올 것이다.

반복문을 돌면서 입력받은 값이 heap[0] 원소보다 크다면 heap[0] 을 pop하고 입력받은 값을 새롭게 heap에 집어넣어 가면서 힙의 크기가 n인 min heap을 유지한다.

그리고 최종적으로 답은 heap의 가장 작은 원소인 heap[0]이 될 것이다.