[Python] boj no. 1655 : 가운데를 말해요

Problem

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

Idea

최대힙과 최소힙을 이용하여 힙 안의 수를 적절하게 이동시켜주면 해당 수들의 중간 값을 구할 수 있을 것이라는 아이디어를 생각했다. 다음과 같이 문제를 해결했다.

1. 반복문마다 주어지는 num값과 최대힙의 rear값을 비교하여 num > rear이면 최소힙에, 아니면 최대힙에 푸쉬하였다.
2. num을 푸쉬한 후, 각 힙의 길이를 비교하여 한 쪽 힙의 원소를 팝하고 다른 한쪽의 힙에 푸쉬함으로써, 최대힙의 길이가 최소힙의 길이보다 1만큼 크도록하였다.
3. 위의 과정을 통해 두 힙의 rear 값은 해당 수들의 중간 혹은 그와 가까운 값이 된다. 최대힙은 rear값을 기준으로 그보다 작은 수들이, 최소힙은 그보다 큰 수들이 원소로 저장된다. 이 때, 최대힙의 길이는 최소힙보다 1만큼 크기 때문에 해당 수들의 중간값은 최대힙의 rear 값이 된다.

Code

import heapq
import sys

input = lambda : int(sys.stdin.readline())

N = input()
minh, maxh = [], []
for _ in range(N):
    num = input()
    
    if len(maxh) == 0:
        heapq.heappush(maxh, -num)
        print(num)
        continue
    
    if num < -maxh[0]:
        heapq.heappush(maxh, -num)
    else:
        heapq.heappush(minh, num)
        
    if len(maxh) < len(minh):
        heapq.heappush(maxh, -heapq.heappop(minh))
    elif len(maxh) > len(minh)+1:
        heapq.heappush(minh, -heapq.heappop(maxh))
        
    print(-maxh[0])

Comment