백준 2696번: 중앙값 구하기 [python]

난이도 : 골드 2
유형 : 우선순위 큐
https://www.acmicpc.net/problem/2696

문제 접근

여러 테스트케이스가 주어지고, 각 테스트마다 수열이 최대 10,000개까지 주어진다.
이 수열을 앞에서부터 하나씩 보면서, 홀수 번째 수를 입력받을 때마다 중앙값을 출력하는 문제이다.

출력은 한 줄에 최대 10개까지만 출력해야한다.

아이디어

최대큐와 최소큐 두개의 큐에 수들을 저장한다.
중앙값을 최대큐인 -minheap 에 넣을 것이다.
그래야 우리가 -minheap[0] 을 통해 바로 중앙값을 가져올 수 있기 때문이다.
최소 큐에 넣으면 중앙값을 바로 가져올 방법이 없다.

1. 입력이 들어올 때, 중앙값(즉 최대 큐의 루트 노드)과의 크기를 비교한다. 중앙값 이하라면 최대 큐에 삽입하고, 그렇지 않다면 최소 큐에 삽입한다.
2. 또한 중앙값의 특성상 최대 큐는 최소 큐와 크기가 같거나, 길이가 1 더 길어야 한다(이렇게 해야 최대 큐의 루트 노드가 항상 중앙값임이 보장된다) 만약 그렇지 않다면, 힙 삽입 및 삭제를 통해 크기를 조정해주도록 한다.

해답 및 풀이

import heapq
import sys
input = sys.stdin.readline

T = int(input())  # 테스트 케이스 수

for _ in range(T):
    m = int(input())  # 이 테스트케이스에서 입력될 숫자 개수
    nums = []

    # 숫자가 여러 줄에 걸쳐서 들어오므로, 전체 m개가 될 때까지 입력 받음
    while len(nums) < m:
        nums += list(map(int, input().split()))  # +=로 이어붙여야 덮어쓰기 안됨

    max_heap = []  # 중앙값 이하를 저장하는 최대 힙 (파이썬은 min-heap이라 음수로 저장)
    min_heap = []  # 중앙값 초과를 저장하는 최소 힙
    result = []    # 홀수 번째 입력마다 저장할 중앙값 리스트

    for i in range(m):
        num = nums[i]

        # max_heap이 비었거나, 현재 수가 중앙값 이하일 경우 max_heap에 저장
        if not max_heap or num <= -max_heap[0]:
            heapq.heappush(max_heap, -num)
        else:
            # 중앙값 초과일 경우 min_heap에 저장
            heapq.heappush(min_heap, num)

        # 두 힙의 크기를 맞춰줌 (max_heap은 항상 min_heap보다 같거나 1개 더 많아야 함)
        if len(max_heap) > len(min_heap) + 1:
            heapq.heappush(min_heap, -heapq.heappop(max_heap))  # max → min 이동
        elif len(min_heap) > len(max_heap):
            heapq.heappush(max_heap, -heapq.heappop(min_heap))  # min → max 이동

        # 홀수 번째 입력일 경우 (0-based index 기준), 중앙값 저장
        if i % 2 == 0:
            result.append(-max_heap[0])  # 중앙값은 항상 max_heap 루트에 있음

    print(len(result))  # 중앙값의 개수 출력
    for i in range(0, len(result), 10):
        # 결과를 10개씩 끊어서 출력
        print(' '.join(map(str, result[i:i+10])))

최대 큐와 최소 큐 두개다 활용하는 우선순위 큐의 화룡점정같은 문제.
잘 이해 한다면 우선순위 큐를 잘 다룰 수 있을 것 같다.