백준 Nr 1655 가운데를 말해요 [Python] 풀이

1. 문제 설정 확인

1. 우리는 N번만큼 정수를 입력받는다

2. 입력받을 때마다 지금까지 입력받은 수들 중 중간값을 출력한다.

2. 문제 풀이

1655에선 특이하게 항상 현재까지 입력 된 수들 중 중간값을 출력해야한다.

물론, 일반적으로 가장 좋은(쉬운) 방법은 for문을 한 번 돌때마다 sort하고(아마 이 경우는 간단한 insertion sort가 좋을 듯 하다.) Length / 2 를 인덱스로 갖는 요소를 출력하면 되지만, 그렇다면 이 문제가 골드 2일 이유가 전혀 없다(…)

여기서 우리는 heap를 사용해 문제를 해결하여야 한다. 그렇다면 heap를 이용해 중간값을 출력하는 방법은 뭐가 있을까?

Heap의 종류는 Max Heap과 Min Heap로 나뉜다. 그렇다면 이 둘을 적절히 섞으면 중간값을 빠르게 구할 수 있지 않을까? 아이디어는 다음과 같다 :

0. Max Heap과 Min Heap 용 array 2개를 만든다.

1. Max Heap과 Min Heap의 길이가 일치하도록 번갈아가며 숫자를 넣어준다.#### 1-1. 단, Min Heap의 경우에는 Min Heap중 가장 큰 숫자가 위로 오게 하기 위헤 숫자에 minus를 붙혀서 넣는다.

2. 단, Min Heap의 가장 위에 있는 숫자가 Max Heap의 가장 위에 있는 숫자보다 크다면, 두 숫자를 바꿔준다.

2-1. 이렇게 하면 Min Heap에는 항상 Max Heap보다 작은 숫자가 오게되고, Min Heap의 가장 꼭대기에는 Max Heap보다 작은 숫자들 중 가장 큰 수가 오게 된다.

2-2. 1번의 조건에 의해 두 heap의 길이는 같으므로, Min Heap[0]은 중간값이 된다.

3. 과정이 끝나면 Min Heap의 가장 위에 있는 숫자를 출력해준다.

3. 해답 코드

Python은 기본적으로 최소 힙만을 제공한다. 따라서 Min Heap의 구현에 있어서는 조금 생각을 하면서 따라와야 한다. 하지만 큰 어려움은 없을것이다. 단지 Max Heap에서 가장 작은 요소보다 작은 요소들을 Min Heap에 넣는데, Min Heap는 최대 힙이라 가장 위에 가장 큰 값이 올라올 뿐이다.

import sys
import heapq

input = sys.stdin.readline
N = int(input())
max_heap = []
min_heap = []

for _ in range(N) :
    num = int(input())
    if len(max_heap) == len(min_heap) :
        heapq.heappush(min_heap, -num)
    else :
        heapq.heappush(max_heap, num)
    if min_heap and -min_heap[0] > max_heap[0] :
        A = heapq.heappop(max_heap)
        B = heapq.heappop(min_heap)
        heapq.heappush(min_heap, -A)
        heapq.heappush(max_heap, -B)
    print(-min_heap[0])

Source

Baekjoon, 2022.07.18
Pycharm