BOJ-1655 가운데를 말해요

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

문제 소개

수가 하나씩 추가될 때마다 지금까지 추가된 수들의 중간값을 출력하는 문제이다. 우선순위 큐를 이용하여 해결하였다. 힙을 사용하기 위해 heapq 모듈을 import하였다.

문제 풀이

먼저 중간값보다 작은 수들이 들어갈 힙을 heapL, 중간값보다 큰 수들이 들어갈 힙을 heapR이라고 선언한다.

heapL은 최대 힙이며, heapR은 최소 힙으로 동작한다. 그리고 왼쪽 힙과 오른쪽 힙을 번갈아 가며 입력된 정수를 push하게 되는데, push 직후 왼쪽 힙의 최대값이 오른쪽 힙의 최소값보다 크다면 그 두 값을 pop한 후 서로 반대 힙으로 push한다.

위 동작을 통해 언제나 heapL의 최대값은 heapR의 최소값보다 작거나 같게 된다.

중간값 결정

인덱스 i를 1부터 n까지로 하여 왼쪽 힙부터 입력을 받게 하면 중간값은 아래와 같이 결정된다.

• n이 홀수인 경우 :
  heapL의 길이가 heapR의 길이보다 1 크며, heapL[0]은 중간값이다.
• n이 짝수인 경우 : heapL과 heapR의 길이는 동일하며, heapL[0]은 문제의 조건에 맞는 수이다.

즉, 언제나 heapL[0]이 중간값이 된다.

코드

import heapq
from sys import stdin
input = stdin.readline

n = int(input())
heapL = []
heapR = []

for i in range(1, n + 1):
    if i % 2 == 0:
        heapq.heappush(heapR, int(input())) # 오른쪽 힙은 최소 힙
    else:
        heapq.heappush(heapL, -int(input())) # 왼쪽 힙은 최대 힙
    try:
        if -heapL[0] > heapR[0]:
            heapq.heappush(heapR, -heapq.heappop(heapL))
            heapq.heappush(heapL, -heapq.heappop(heapR))
    except IndexError:
        pass
    # print(f"heapL : {heapL}")
    # print(f"heapR : {heapR}")
    print(-heapL[0])

try except문을 사용한 이유는 i가 1일 때, heapR이 빈 리스트인 상황을 넘기기 위해서이다.

주석처리한 2개의 print문으로는 매 반복마다 두 힙의 내용을 확인할 수 있다.