[BOJ] N번째 큰 수 2075 - Python

문제

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번째 큰 수를 출력한다.

해결 방법

• 입력된 수 중 N번째 큰 수를 찾으면 되는 문제
• list로 받아서 정렬 후 n번째 큰 수를 받으면 되긴 하지만, 메모리 초과 발생
• n의 최댓값이 1500이기 때문에, 최악의 경우 1500*1500개의 데이터를 가진 리스트를 생성해야 함
• heap을 사용해 문제 해결
• heap에 n개의 수만 저장하고, 만약 n보다 많아진다면, 제일 작은 수를 빼고 현재 수를 heap에 추가

⌛ 메모리 초과

def solution():
    n = int(input())
    digits = []
    for _ in range(n):
        for i in map(int, input().split()):
            digits.append(i)
    digits.sort(reverse=True)
    print(digits[n-1])

CODE

import heapq

def solution():
    n = int(input())
    heap = []
    for _ in range(n):
        for x in map(int, input().split()):
            if len(heap) >= n:
                heapq.heappushpop(heap, x)
            else:
                heapq.heappush(heap, x)
    print(heap[0])

시간복잡도