BOJ 1655: 가운데를 말해요 [Python 3]

골드 2의 난이도를 자랑하는 '가운데를 말해요'라는 문제이다.

문제의 내용은 간단하다.

모든 input이 들어올 때 마다 아직까지 들어온 input들의 중앙값을 출력해주면 된다.

하지만 정말 큰 문제는 해당 문제의 시간 제한이 0.1초이라는 점이다.

그렇기 때문에, 단순히 input이 들어올 때 마다 sort를 해서 중앙값을 출력해주면 너무너무 느려진다. 

그래서 최소 및 최대 힙을 사용해서 이 문제를 풀어보았다.

 

import heapq
import sys
input=sys.stdin.readline
maxHeap=[]
minHeap=[]
n=int(input())
for i in range(n):
    num=int(input())
    if len(maxHeap)==len(minHeap):
        heapq.heappush(minHeap, -1*num)
    else:
        heapq.heappush(maxHeap, num)
    if maxHeap and -minHeap[0]>maxHeap[0]:
        a=heapq.heappop(minHeap)
        b=heapq.heappop(maxHeap)
        heapq.heappush(minHeap,-b)
        heapq.heappush(maxHeap,-a)
    print(-minHeap[0])

중간값보다 작은 최대 힙, 중간값보다 큰 최소 힙을 만든다.

이렇게 되면, 중간 값은 최소 힙의 head가 된다. 

만약 개수가 짝수개라면, 작은 수를 대답해야 하기 때문에, 먼저 최소 힙부터 채운다.