[Algorithm] 백준 1655번: 가운데를 말해요 | 최대 힙, 최소 힙 응용

문제 출처: https://www.acmicpc.net/problem/1655

---

1. 결과

메모리 35032KB, 시간 272ms

2. 풀이

이 문제는 N개의 숫자가 각각 입력될 때마다, 입력된 숫자들 중에서 중간값을 출력해야 한다. 평균값이나 누적 합이 아니기 때문에 각각의 값이 살아서 존재해야 하고, 값이 새로 입력될 때마다 정렬이 수행되어야 중간값을 구할 수 있다. 하지만 값이 추가될 때마다 정렬(sort)이 발생하고, 이런 과정이 N번 반복된다면 시간 복잡도는 커질 수밖에 없다. 이를 더 효율적인 방법으로 풀기 위해서는 최대 힙과 최소 힙을 함께 응용하는 방법이 있다.

우선, 최대 힙과 최소 힙이 의미하는 바를 이해해야 하는데, 최대 힙은 [0] 번째 원소부터 내림차순으로 저장 혹은 출력되는 자료구조이다. 최소 힙은 [0] 번째 원소부터 오름차순으로 저장 혹은 출력되는 자료구조이다. 이를 정렬된 수열에 적용했을 때, 중간값을 기준으로 왼쪽으로 바라보면 최대 힙, 오른쪽 방향으로 바라보면 최소 힙이 된다. 예를 들어 아래 그림 1처럼 이미 정렬된 배열이 있다고 가정했을 때, 중간값을 기준으로 각각 그 아래의 그림 2, 그림 3처럼 인덱스를 갖는 최대 힙과 최소 힙으로 나눌 수 있다는 것이다.

그림1. 이미 정렬된 배열

그림2. 그림1에서 중간값을 기준으로 왼쪽(중간값 포함)으로 바라본 최대 힙

그림3. 그림1에서 중간값을 기준으로 오른쪽으로 바라본 최소 힙

단, [0]번째 인덱스 원소 이후의 값들은 실제 힙큐로 구현한 우선순위 큐와 배열이 다를 수 있다. heapq의 heappush로 값을 저장할 때, 실제 우선순위가 높은 순서 그대로 값을 저장하지는 않기 때문이다. 대신 [0] 번째 값은 가장 우선순위가 높은 원소가 맞기 때문에 [0] 번째 값은 확실하다. 위처럼 오른쪽(최소 힙)과 왼쪽(최대 힙)의 배열 길이를 같거나, 왼쪽이 1만큼 더 큰 상태로 유지하면서, heappush와 heappop으로 값을 저장하고, 옮기며 매 값이 입력될 때마다 최대 힙의 [0] 번째 값을 출력하면 정답을 도출할 수 있다. 이를 구현한 설명과 코드는 아래와 같다.

import sys, heapq

# 0. 중간값은 숫자를 나열했을 때(예시: 1, 2, 3, 4, 5), 반을 기준(홀수일 경우 최대 힙에 +1)으로 최대 힙과 최소 힙으로 분리하고,
# 최대 힙(3, 2, 1)의 [0]번째 값을 중간값으로 정의한다.
# 1. 최대 힙, 최소 힙을 생성
# 2-1. 숫자 하나를 입력 받았을 때, 최대 힙이 비어있거나, 최대 힙의 [0]번째보다 작으면 최대 힙에 push, 
#      push 후에 '최대 힙의 길이 > 최소 힙의 길이 + 1'이면 최대 힙에서 pop한 값을 최소 힙에 push
# 2-2. 아닌 경우에는 최소 힙에 push
#      push 후에 '최대 힙의 길이 < 최소 힙의 길이'라면 최소 힙에서 pop한 값을 최대 힙에 push
N = int(sys.stdin.readline())
max_hq, min_hq = [], []

for _ in range(N):
    num = int(sys.stdin.readline())
    if not max_hq or num < -max_hq[0]:
        heapq.heappush(max_hq, -num)
        if len(max_hq) > len(min_hq) + 1: heapq.heappush(min_hq, -heapq.heappop(max_hq))
    else:
        heapq.heappush(min_hq, num)
        if len(max_hq) < len(min_hq): heapq.heappush(max_hq, -heapq.heappop(min_hq))
    print(-max_hq[0])