백준 1655 가운데를 말해요, 파이썬

문제

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

문제를 읽고 heapq를 사용해야겠다는 생각은 들었지만 코드를 어떻게 짜야할지 모르겠어서 검색을 했다...

풀이

힙을 두 개 써서 n을 분배한다.

왼쪽 힙은 최대 힙으로 정렬하고, 오른쪽 힙은 최소 힙으로 정렬하면 왼쪽 힙의 첫번째 요소는 항상 중앙값이 된다.

1. 왼쪽 힙과 오른쪽 힙의 길이가 같으면 (x * -1)을 왼쪽 힙에 삽입한다.
   -> 그렇지 않으면 오른쪽 힙에 삽입한다.
2. 왼쪽 힙과 오른쪽 힙에 요소가 존재하고, 왼쪽 힙의 (첫번째 수 * -1)가 오른쪽 첫번째 요소보다 클 때 -> 왼쪽 힙의 첫번째 요소와 오른쪽 힙의 첫번째 요소를 바꿔준다. (-1을 곱하고 바꿔준다.)
3. 왼쪽 힙의 (첫번째 수 * -1)을 출력한다.

코드

import sys
input = sys.stdin.readline
import heapq

n = int(input())
left_heap, right_heap = [], []

for i in range(n):
    x = int(input())

    if len(left_heap) == len(right_heap):
        heapq.heappush(left_heap, -x)
    else:
        heapq.heappush(right_heap, x)

    if len(left_heap) >= 1 and len(right_heap) >= 1 and left_heap[0] * -1 > right_heap[0]:
        max_num = heapq.heappop(left_heap) * -1
        min_num = heapq.heappop(right_heap)
        
        heapq.heappush(left_heap, min_num * -1)
        heapq.heappush(right_heap, max_num)

    print(left_heap[0] * -1)