[백준 1655번] 가운데를 말해요

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

import sys
import heapq
input = sys.stdin.readline

# 정수를 하나씩 외칠때마다 지금까지 말한 수 중에서 중간값을 출력
# 짝수개라면 중간에 있는 두 수 중에서 작은 수

n = int(input())

heap1 = [] # 중간값을 포함한 이전 값 최대힙
heap2 = [] # 중간값 이후의 값 최소힙

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

    if len(heap1) == len(heap2):
        heapq.heappush(heap1,-num)
    else:
        heapq.heappush(heap2,num)

    if heap2 and heap2[0] < -heap1[0]:
        lvalue = heapq.heappop(heap1)
        rvalue = heapq.heappop(heap2)
        heapq.heappush(heap1,-rvalue)
        heapq.heappush(heap2,-lvalue)
    
    print(-heap1[0])

 
이번 문제는 n개의 정수를 입력받을 때마다 매번 입력받은 값 중에서 중간 값을 출력해야 하는 문제이다.
 
문제의 입력 조건(1 <= n < 10^5)과 제한시간(0.6초)을 보자마자 n개의 입력을 받을 때마다 정렬(O(NlogN))시키게 되면 시간초과가 발생한다는 것을 인지할 수 있었다.
 
그래서 최대와 최소와 관련된 자료구조 Heap을 활용하기 위해 heapq를 사용해야겠다고 생각하긴 했는데, heapq는 부모와 자식 노드간의 대소관계를 통해 최대와 최소값을 보장해줄 뿐 형제 노드 간에는 대소관계가 구분되지 않기 때문에 단순히 heapq 내부의 값들에서 중간 인덱스를 출력하면 엉뚱한 값이 나오게 된다.
 
따라서 이 문제는 다음과 같이 입력받은 모든 값을 2개의 '우선순위 큐'로 나눠서 풀이해야 한다.
 

1. 중간값과 중간값보다 작은 값 중 가장 큰 값을 반환할 최대 힙(heap1)과 중간값보다 큰 값 중 가장 작은 값을 반환하는 최소 힙(heap2)을 선언
2. 입력받은 값의 개수가 짝수인 경우(len(heap1) == len(heap2))에는 중간의 2개의 수 중 작은 수를 출력해야 하므로 왼쪽 최대 힙(heap1)에 원소를 넣어준다.
3. 입력받은 값의 개수가 홀수인 경우에는 좌우 힙들의 길이를 맞춰주기 위해 오른쪽 최소 힙(heap2)에 원소를 넣어준다.
4. 원소를 추가한 이후에는 왼쪽 최대 힙과 오른쪽 최소 힙 값의 오름차순 정렬을 유지하기 위해 왼쪽 최대 힙의 최댓값(-heap1[0])과 오른쪽 최소 힙의 최솟값(heap2[0])을 비교한 뒤, 값이 올바르지 않으면 서로 교환해준다.
5. 왼쪽 최대 힙의 최댓값(-heap1[0])을 출력해준다.

우선순위 큐와 heapq를 활용하는 방법을 익힐 수 있었던 문제였다 :)