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

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

 

# 백준 1655 가운데를 말해요
# 중간값을 실시간으로 구하는 문제
# 최대 힙(left)과 최소 힙(right)을 사용하여 중간값을 효율적으로 관리

import sys, heapq
input = sys.stdin.readline
N = int(input())  # 수의 개수 입력

lq = []  # left heap (최대 힙): 중간값 이하를 저장 (파이썬에서는 -를 붙여서 최대 힙처럼 사용)
rq = []  # right heap (최소 힙): 중간값 초과를 저장

lc = 0  # left heap에 들어있는 원소 개수
rc = 0  # right heap에 들어있는 원소 개수

ans = []  # 출력할 정답 리스트

for _ in range(N):
    num = int(input())

    if lc <= rc:
        # 최소 힙(rq)에 먼저 넣고, 그 중 최소값을 최대 힙(lq)으로 이동
        # 이렇게 하면 lq가 rq보다 항상 크거나 같아짐 → 중간값은 lq의 top
        if rq:
            heapq.heappush(rq, num)
            heapq.heappush(lq, -heapq.heappop(rq))
        else:
            heapq.heappush(lq, -num)
        lc += 1
    else:
        # 최대 힙(lq)에 먼저 넣고, 그 중 최대값을 최소 힙(rq)으로 이동
        heapq.heappush(lq, -num)
        heapq.heappush(rq, -heapq.heappop(lq))
        rc += 1

    # 항상 최대 힙의 top이 중간값
    ans.append(-lq[0])

# 결과 출력
for a in ans:
    print(a)

우선순위 큐 두개를 이용하여 해결하였다.

 

왼쪽에 있는 우선순위큐는 높은수를

오른쪽에 있는 우선순위큐는 낮은수를 반환하도록 하였다.

 

이에따라 새로운 수가 삽입될때마다 수를 적절히 왼쪽/오른쪽으로 이동시켜 가운데 수를 바로 찾을 수 있도록 구현하였다.