[백준/Baekjoon] 2075번: N번째 큰 수

1. 문제

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

2. 입력

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

3. 출력

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

4. 예제 입력 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

5. 예제 출력 1

35

6. 풀이

힙(Heap) 은 힙의 특성(최소 힙(Min Heap) 에서는 부모가 항상 자식보다 작거나 같다) 을 만족하는 거의 완전한 트리인 특수한 트리 기반의 자료구조다. 이를 구현하기 위해 heapq 모듈 을 사용한다.

일반적인 방법으로는 메모리 초과가 발생하므로, N개의 원소를 갖는 최소 힙을 이용하여 문제를 풀면 된다.

기본적인 원리는 최소 힙의 첫 번째 원소가 가장 작은 값이라는 점을 이용한다. 최소 힙 내의 모든 값들을 모든 원소 중 상위 N개를 가지도록 구현하는 것이다. 즉, 입력한 값이 최소 힙에서 가장 작은 값보다 큰 값이라면 이 값을 힙에 삽입하고, 가장 작은 값은 꺼낸다. 이러한 과정을 반복하면 최소 힙은 상위 N개의 원소를 가지며, 이 중 첫 번째 원소는 최소 힙의 원리에 의해 N번째 큰 수가 된다.

data = map(int, input().split())
for d in data:
    heapq.heappush(heap, d)  # 첫 입력은 전부 최소 힙에 삽입

answer = heap[0]  # 최소 힙의 첫 번째 원소가 N번째 큰 수

for _ in range(n - 1):
    data = map(int, input().split())
    for d in data:
        if d > answer:  # 입력한 값의 원소가 최소 힙의 첫 번째 원소보다 큰 값이라면
            heapq.heappush(heap, d)  # 입력한 값을 최소 힙에 삽입
            heapq.heappop(heap)  # 첫 번째 원소는 꺼낸다
            answer = heap[0]  # N번째 큰 수를 다시 최소 힙의 첫 번째 원소로 지정

7. 코드

import heapq

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

data = map(int, input().split())
for d in data:
    heapq.heappush(heap, d)

answer = heap[0]

for _ in range(n - 1):
    data = map(int, input().split())
    for d in data:
        if d > answer:
            heapq.heappush(heap, d)
            heapq.heappop(heap)
            answer = heap[0]

print(answer)