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

kimminjunnn·2025년 12월 14일

알고리즘

목록 보기
264/311

문제 출처 : https://www.acmicpc.net/problem/2696
난이도 : 골드 2


문제 파악

숫자가 M개 주어질 때, 1번째, 3번째, 5번째… (즉 지금까지 읽은 개수가 홀수일 때) 마다 “지금까지의 중앙값”을 출력해야 한다.

1개 읽으면 그 값이 중앙값
3개 읽으면 정렬했을 때 가운데 값이 중앙값
5개 읽으면 정렬했을 때 가운데 값이 중앙값
.
.

출력은:

중앙값 개수 ((M+1)//2) 먼저 출력,
중앙값들을 한 줄에 10개씩 끊어서 출력해야 한다.


줄바꿈 된 입력을 같은 리스트에 저장해야 하기 때문에 파이썬에 extend 함수를 알면 좋을 것 같다.
extend는 리스트에 다른 리스트(또는 여러 원소)를 펼쳐서 붙이는 함수이다.


핵심 아이디어

매번 “지금까지 값들 정렬해서 중앙값”을 구하면, M이 최대 9999개 에다가 테스트케이스 도 여러개 일 수 있어 시간초과가 날 것 같다.

중앙값 문제는 힙 2개를 사용해 풀 수 있다.


힙 구성

max_heap에는 “작은 절반”이 들어가게 한다. ex) 1,2,3
min_heap에는 “큰 절반”이 들어가게 한다. ex) 4,5

크기는 항상 len(max_heap) == len(min_heap) 또는 len(max_heap) == len(min_heap) + 1 유지한다.

즉, 홀수 개일 때는 max_heap이 1개 더 많게 맞춘다.

그러면 홀수 개일 때 중앙값은 항상 -max_heap[0] -> 3

홀수 개에서 작은 절반이 하나 더 많으면, “가운데 값”이 작은 절반의 최대값이 되기 때문이다.


import sys
import heapq

input = sys.stdin.readline

T = int(input().strip())
out = []

for _ in range(T):
    M = int(input().strip())

    # --- 입력: M개가 모일 때까지 계속 읽기 (줄에 10개씩 끊겨도 OK) ---
    nums = []
    while len(nums) < M:
        nums.extend(map(int, input().split()))

    max_heap = []  # 작은 절반 (최대힙 역할: 음수로 저장)
    min_heap = []  # 큰 절반 (최소힙)
    medians = []

    for i, x in enumerate(nums, 1):  # i = 지금까지 읽은 개수
        # 1) 일단 넣기
        if not max_heap or x <= -max_heap[0]:
            heapq.heappush(max_heap, -x)
        else:
            heapq.heappush(min_heap, x)

        # 2) 밸런스 맞추기
        if len(max_heap) < len(min_heap):
            heapq.heappush(max_heap, -heapq.heappop(min_heap))
        elif len(max_heap) > len(min_heap) + 1:
            heapq.heappush(min_heap, -heapq.heappop(max_heap))

        # 3) 홀수번째마다 중앙값 기록
        if i % 2 == 1:
            medians.append(-max_heap[0])

    # --- 출력: 중앙값 개수 + 중앙값들을 10개씩 끊어서 ---
    out.append(str(len(medians)))
    for j in range(0, len(medians), 10):
        out.append(" ".join(map(str, medians[j:j + 10])))

print("\n".join(out))
profile
Frontend Engineers

0개의 댓글