[우선순위 큐, 힙] 백준 1658번 가운데를 말해요 with Python3 ★

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

1655번: 가운데를 말해요

 

문제

 

정답 풀이

import heapq
import sys

# 중앙값 기준으로 작은 값 = left, 큰 값 = right

left, right = [], []
n = int(sys.stdin.readline())

for _ in range(n):
    num = int(sys.stdin.readline())

    # 최소 힙과 최대 힙에 원소를 번갈아가며 push/ left부터이긴함
    if len(left) == len(right):
        # max_heap
        heapq.heappush(left, (-num, num))  # 최대힙 기본은 최소이기에 역순으로

    else:
        # min heap.
        heapq.heappush(right, (num, num))

    # 오른쪽 값에 원소가 있으면서
    # 왼쪽 값이 오른쪽 값보다 큰 경우:
    # -> left원소보다 right원소가 더 커야하는 조건에 위배

    # 그렇기에 원소 위치 교체
    if right and left[0][1] > right[0][1]:
        left_val = heapq.heappop(left)[1]
        right_val = heapq.heappop(right)[1]
        heapq.heappush(right, (left_val, left_val))
        heapq.heappush(left, (-right_val, right_val))

    # 결과적으로 왼쪽 최대힙의 첫번째 원소가 항상 중앙값이 된다.
    print(left[0][1])

출처:https://inspirit941.tistory.com/200

 

해석

최소 힙과 최대 힙을 사용한다는 점에서 생각보다 간단한 내용이었다. 굳이 코드 까진 아니고 구상만 보고 직접 구현해볼 수도 있었을거 같은데 아쉽다.