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

 

 

from heapq import *
n = int(input())
heap = [] 
for _ in range(n):
    for i in list(map(int, input().split())): # 리스트를 돌면서 원소를 힙에 넣어줌
        if len(heap) < n: # 만약 힙의 요소가 n개 미만이라면 i를 힙에 추가
            heappush(heap, i)
        else: # 힙의 요소가 n개 이상일때 
            if heap[0] < i: # 만약 최소값이 i보다 작다면
                heappop(heap) # 현재 힙의 최소값을 제거
                heappush(heap, i) # 현재 요소를 힙에 추가 
print(heap[0])

 

처음에는 최대 힙을 구하고 n번째가 될때까지 pop을 반복하는 코드를 짰었다. 하지만 결과는.....^^

입력값이 n^2이기때문에 확실히 좋은 방법이 아니었고 모든 입력값을 리스트에 넣은채 최대 힙으로 변환했기때문에 메모리에 문제가 생겼을 수 있을것 같았다. 그래서 입력값을 줄일수 있는 방법에 대해서 고민하게 됐다. 시간을 들였는데도 잘 안되서 결국 검색을 했고, 힙의 길이를 n개로 제한하는 방법을 알게 되었다. 

코드를 보면 리스트를 돌면서 힙의 길이가 n이 될때까지 숫자를 넣어준다. 그리고 힙의 길이가 n 이상이 되었을때부터 현재 요소(i)와 힙의 최소값을 비교한다. 코드를 처음 봤을때는 n번째로 큰 값을 찾는건데 왜 최소값과 비교하는지 이해가 되지 않았다. 하지만 다시 보면 최소값과 비교하는 이유는 길이 n의 힙에서 최소값들은 제거하고 최대값만 남기기 위해서이다. n번째로 큰 수는 바꿔말하면 최대값 기준으로 정렬했을때 n번째로 작은 수이기도 하다. 그래서 최소값과 현재 숫자(i)를 비교하는 것이다. 

힙의 길이가 n 이상일때 현재 숫자가 힙의 최솟값 보다 작다면 무시하고 현재 숫자가 힙의 최솟값보다 크다면 그 값을 제거하고 현재 숫자를 힙에 추가한다. 그러면 최종적으로는 n개의 요소가 들어있는 최소힙에는 입력값들 중 가장 큰 순서대로 n개가 모인다. 그중에서의 최소값이 n번째 최대값이므로 마지막에는 힙의 0번째 요소를 리턴한다. 

 

이렇게 구현하는 방법도 있고, 나중에 알게 된 내용이지만 파이썬 heapq 라이브러리에 이미 n번째 큰 값을 구하는 함수가 있었다...ㅋㅋㅋㅋㅋㅋ heapq.nlargest()는 힙에서 n개의 가장 큰 요소로 구성된 리스트를 반환하고, 반대로 heapq.nsmallest()는 힙에서 n개의 가장 작은 요소로 구성된 리스트를 반환한다. 즉 힙에서 큰/작은 순서대로 n개의 요소가 들어있는 리스트를 반환해준다.

from heapq import *

heap = [1, 2, 3, 4, 5, 6, 7, 8, 9]
heapify(heap)

print(nlargest(3, heap)) # [9, 8, 7]
print(nsmallest(3, heap)) # [1, 2, 3]

 

 

 

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

2075번: N번째 큰 수