백준 BOJ 2075 N번째 큰 수

문제

코드

import sys
import heapq

input = sys.stdin.readline

 if __name__ == "__main__":
    n = int(input())
    metrix = [[*map(int, input().split())] for _ in range(n)]
    heap = [*metrix[0]]
    heapq.heapify(heap)
    
    for i in range(1, n):
        for j in range(n):
            if metrix[i][j] > heap[0]:
                heapq.heappop(heap)
                heapq.heappush(heap, metrix[i][j])
    
    print(heap[0])

 

항상 상위 N개만 유지해야 하고, 그 N개 중에서 가장 작은 값이 전체에서 N번째 큰 값이므로 heap을 사용한다.

최소힙을 써서 가장 작은 값인 루트 값을 리턴한다.

 

 

 

위와 같이 풀었으나 챗선생이 아래와 같이 조언해주었다.

import sys
import heapq

input = sys.stdin.readline

n = int(input())

heap = list(map(int, input().split()))
heapq.heapify(heap)

for _ in range(n - 1):
    nums = map(int, input().split())

    for num in nums:
        if num > heap[0]:
            heapq.heapreplace(heap, num)

print(heap[0])

 

굳이 metrix를 전체 받지 말고 메모리 아끼고 효율적인 처리를 하는 겸 첫번째 줄은 heap에 바로 넣고 이후 입력을 반복으로 받아서 처리하는 쪽으로 제안했다.

이게 더 좋은 것 같다.

 

N by N 행렬에서 N번째 큰 수를 찾는 것인데, 일차적으로 생각하면 heap에 다 때려넣고 N-1번째 꺼내면 되는거 아닌가? 싶지만 이는 메모리 초과가 나올 수 밖에 없다(n^2).

 

따라서 heap을 N 사이즈로 두고, 새로 들어온 입력이 heap보다 크면 heap에 넣고, 그렇지 않으면 패스하는 방식으로 코드를 작성한다.

이때, heapify 는 리스트를 힙으로 전환시켜주는 함수이다.

 

heaprelpace는 나도 처음 알았는데, 루트(pop) 제거 + 새로운 값 push 이다.

heapq.heapreplace(heap, x)

 

즉, 아래 두개를 함께하는 함수이다.

heapq.heappop(heap)
heapq.heappush(heap, x)

 

동작을 실제로 보면

heap = [3, 5, 7]  # 최소힙

heapq.heapreplace(heap, 6)

1. 3 제거
2. 6 삽입
→ [5, 6, 7]

이렇게 된다.

 

그래서 조건없이 쓰면 안된다.

if num > heap[0]:
    heapq.heapreplace(heap, num)

 

 

비슷하게 heappushpop 이라는 함수가 있는데, 

heapq.heappushpop(heap, num)

 

이건

push 먼저 → 그 다음 pop

 

순서대로 이므로 결과가 다를 수 있다.

 

 

 

TOP-K, K번째 작은 값

K번째 큰 값

큰 값 K개만 남기기 -> 그 중 가장 작은 값이 K번째 큰 값이다.

따라서 min heap을 사용하며, heap[0]은 K번째로 큰 값이 된다.

import heapq

heap = []

for num in data:
    if len(heap) < k:
        heapq.heappush(heap, num)
    else:
        if num > heap[0]:
            heapq.heapreplace(heap, num)

print(heap[0])

 

 

K번째 작은 값

작은 값 K만 남기기 -> 그 중 가장 큰 값이 K번째 작은 값이다.

따라서 max heap을 사용하며 음수로 처리한다.

import heapq

heap = []

for num in data:
    if len(heap) < k:
        heapq.heappush(heap, -num)
    else:
        if num < -heap[0]:
            heapq.heapreplace(heap, -num)

print(-heap[0])