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

 

 

이 문제의 핵심은 최소힙을 사용하는 것이었다. 사실 처음엔 배열이나 스택을 활용하려고 했었는데 너무 모르겠어서 구글링 해보니 다들 최소힙으로 푸셨음..ㅎ

즉 최소힙의 특성을 기반으로 상위 N개의 큰 숫자만 유지하는 방식을 구현해야 하는 문제였다.

 

최소힙의 개념부터 복습해보자면,

최소 힙은 root node가 항상 가장 작은 값을 가지는 이진 트리이다.

파이썬에는 이를 효율적으로 구현할 수 있는 자료구조로 heapq 라이브러리가 이미 존재하는데, 노드 삽입을 위한 heappush와 제거를 위한 heappop 메소드(가장 작은 값인 0번째 원소 제거)가 존재하고 각각 O(log k)의 시간 복잡도를 가진다.

동작 방식은 힙에 새로운 숫자를 삽입할 때 내부적으로 힙 구조를 유지하도록 자동으로 정렬이 되는데, 가장 작은 값은 항상 heap[0]에 위치하게 된다.

따라서 최소 힙을 계속 크기 N으로 유지하면, 힙의 루트에는 입력된 값 중 N번째로 큰 값이 남게된다.

 

디버깅 출력 예시와 아래 코드를 보면서 이해해보자.

 

import sys
import heapq
N = int(input())
heap = [] # 최소 힙
for _ in range(N):
    array = map(int, sys.stdin.readline().split())
    for num in array:
        heapq.heappush(heap, num)
        if len(heap) > N:
            heapq.heappop(heap)

print(heap[0])

 

우선 NxN 배열을 입력받는데, 효율성을 위해 N개의 숫자로 구성된 한 줄(행)씩 입력 받고 정렬하는 과정을 반복한다.

이를 위해 for _ in range(N) 내부에 array에 들어갈 숫자들의 집합을 입력받았고, 내부의 또 다른 중첩 for문을 통해 해당 줄의 숫자들(num)을 차례대로 heappush한다.

 

이때 if len(heap) > N 조건을 삽입하여 heap에는 항상 N개의 원소들로만 구성이 되도록 유지하였고, 만약 하나가 추가되어 heap의 원소가 N+1개이 되는 순간에는 해당 if문을 실행하여 heappop() 메소드로 heap의 0번째 원소(가장 작은 값)를 제거하도록 구현하였다.

 

이렇게 하면 반복문을 실행하면서 heap의 내부에는 항상 N개의 원소들로만 유지가 되고, 계속 숫자들이 추가될 때마다 가장 작은 원소들을 제거함으로써 가장 마지막에는 큰 원소들 N개만 heap에 남아있게 되는 것이다.

 

예제 실행 결과

 

위 예제를 디버깅을 통해 중간 단계 (heap의 상태)를 확인해보면 다음과 같다.

 

 

위에서 볼 수 있듯이 NxN 배열을 N개의 숫자로 이루어진 한 줄씩, N 번 입력하면서 각각 노드가 추가될 때마다 heappop()이 실행되어 heap의 내부 원소 개수는 N개로 유지되도록하고 있는 모습을 볼 수 있다!