boj1655 - 가운데를 말해요

문제

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

코드

import sys
input = sys.stdin.readline
import heapq # 우선순위큐를 이용하여 빠르게 중앙 값 찾기

n = int(input())
maxheap, minheap = [], [] # maxheap의 모든 원소는 minheap의 모든 원소보다 작아야함
for i in range(n):
    if len(maxheap) == len(minheap):
        heapq.heappush(maxheap, -int(input()))
    else:
        heapq.heappush(minheap, int(input()))
    
    if not minheap:
        print(-maxheap[0])
        continue

    if -maxheap[0] > minheap[0]: 
        # maxheap에서 제일 큰 원소가 minheap의 제일 작은 원소보다 크다면 
        # 두 원소를 바꾸어 넣기
        max_tmp = -heapq.heappop(maxheap)
        min_tmp = heapq.heappop(minheap)
        heapq.heappush(minheap, max_tmp)
        heapq.heappush(maxheap, -min_tmp)
    print(-maxheap[0])

풀이

우선순위 큐를 이용하여 빠르게 중앙값을 찾는 문제
우선순위 큐 A, B가 있다고 가정,

조건1 : len(A) == len(B) 또는 len(A) = len(B)+1 임.
조건2 : A의 모든 원소가 B의 모든 원소보다 작으면 A와 B의 전체 원소에서 중앙 값은 A에서 가장 큰값 --> 이 아이디어를 이용하여 구현

입력되는 숫자를 maxheap과 minheap 에 번갈아가며 push하여 조건 1을 만족시킴
숫자 push후 매번 maxheap에서 제일 큰 값(코드에서는 음수값이므로 인덱스 0에 위치한 값)과 minheap에서 제일 작은 값을 비교하여 조건 2를 만족하도록 함

이후 max heap에서 가장 큰값을 출력하면 중앙값