코드

import heapq
n = int(input())

# 중간 값 기준 더 작은 값 (최대힙)
# max_heap의 첫 번째 요소는 최대 힙에서 가장 큰값
max_heap = []

# 중간 값 기준 더 큰 값 (최소힙)
min_heap = []

# 최종 정답을 담을 리스트
answer = []

for _ in range(n):
    num = int(input())
    
    # 두 힙의 길이가 같다면 => 백준이가 외친 수의 개수가 0이거나 짝수
    if len(max_heap) == len(min_heap):
        
        # max_heap(첫 번째 힙, 최대힙)에 넣어줌
        # Python은 기본적으로 최소힙이 구현
        # 최대힙으로 맞추기 위해 num을 음수로 넣어줌
        heapq.heappush(max_heap, (-num))
    else:
        heapq.heappush(min_heap, (num))
        
    # 조건1. min-heap이 null이 아님
    # => max_heap부터 채우므로 max_heap도 당연히 null이 아님
    # => 따라서 min_heap의 null 여부만 체크해줘도 됨
    
    # 조건2. 최대힙에 들어가는 원소의 크기가 더 크면 원소를 바꿔줌
    # => 중간 값보다 큰 원소가 min_heap에 들어감
    # => 항상 중간 값이 max_haep에 있게 만들지 못 함
    if min_heap and -max_heap[0] > min_heap[0]:
        
        # 이를 보완하기 위해
        # 두 힙의 첫 원소를 교체하여 원하는 로직으로 다시 맞춰줌
        _max = heapq.heappop(max_heap)
        _min = heapq.heappop(min_heap)
        
        heapq.heappush(max_heap, (-_min))
        heapq.heappush(min_heap, (-_max))
    
    # 아래와 같이 넣어주면 짝수 개인 경우 작은 수가 들어감
    answer.append(-max_heap[0])
    
for i in answer:
    print(i)

문제 해설

파이썬의 heapq를 이용하여 주어지는 숫자들의 중간 값을 계속해서 출력하는 문제.

처음에 문제를 접근하여 풀 때, 하나의 힙만을 이용하여 시도하였으나, 난관에 봉착하여 결국 구글링을 통해 해법을 찾아보았다. 최대힙과 최소힙, 두 개의 힙을 구현하여 최대힙의 첫 요소에 항상 중간값이 들어올 수 있도록 구현을 하면 생각보다 쉽게 문제를 풀 수 있었다.

파이썬은 기본적으로 heapq가 최소힙을 지원하므로, 최대힙을 만드려면 주어지는 숫자를 음수로 힙에 push하고 pop할 때도 -를 붙여 출력하면 최대힙과 동일한 기능을 구현할 수 있다.

이 문제의 경우 어려운 부분이 있는데, 바로

# 조건1. min-heap이 null이 아님
    # => max_heap부터 채우므로 max_heap도 당연히 null이 아님
    # => 따라서 min_heap의 null 여부만 체크해줘도 됨
    
    # 조건2. 최대힙에 들어가는 원소의 크기가 더 크면 원소를 바꿔줌
    # => 중간 값보다 큰 원소가 min_heap에 들어감
    # => 항상 중간 값이 max_haep에 있게 만들지 못 함
    if min_heap and -max_heap[0] > min_heap[0]:
        
        # 이를 보완하기 위해
        # 두 힙의 첫 원소를 교체하여 원하는 로직으로 다시 맞춰줌
        _max = heapq.heappop(max_heap)
        _min = heapq.heappop(min_heap)
        
        heapq.heappush(max_heap, (-_min))
        heapq.heappush(min_heap, (-_max))

위의 코드이다.

우리가 선택한 중간 값을 담는 힙은 max_heap(최대힙)이다. 허나 만약 max_heap의 첫 번째 요소가 min_heap의 첫 번째 요소보다 크게 입력되는 경우, max_heap의 첫 번째 요소가 항상 중간 값으로 되게 하기 위한 로직이 망가진다. 따라서 이를 보완하기 위해 위의 경우 두 힙의 첫 원소를 서로 교체해주면 결국 다시 max_heap의 첫 요소로 항상 중간 값을 가지는 효과를 볼 수 있다.

댓글남기기

참고