2696: 중앙값 구하기

풀이 시간

• 16m
• 우선순위 큐로 중앙값을 구하는 알고리즘은 이전에 "1655: 가운데를 말해요"에서 풀어본 적이 있었다

구현 방식

• max_heap에는 중앙값보다 작은 값이, min_heap에는 중앙값보다 큰 값이 들어간다

• 두 heap의 길이를 맞춰가면서 값을 넣어주되, 길이가 같은 경우에는 max_heap에 값을 넣어준다
  -> 추가적으로 확인해줘야할 부분은, 만약 -max_heap[0] > min_heap[0]인 경우 (min_heap에 max_heap 보다 작은 값이 들어가는 경우)가 발생한다면 이 둘을 swap 해줘야한다는 점

코드

import sys
import heapq

T = int(sys.stdin.readline()[:-1])
for t in range(T):
    M = int(sys.stdin.readline()[:-1])
    sequence = []
    for _ in range(M//10 + 1):
        sequence.extend(list(map(int, sys.stdin.readline()[:-1].split())))

    max_heap = []; min_heap = []; result = []
    for m in range(M):

        if len(max_heap) == len(min_heap):
            heapq.heappush(max_heap, -sequence[m])
        else:
            heapq.heappush(min_heap, sequence[m])
        
        if min_heap:
            if -max_heap[0] > min_heap[0]:
                low = -heapq.heappop(max_heap); high = heapq.heappop(min_heap)
                heapq.heappush(max_heap, -high)
                heapq.heappush(min_heap, low)
        
        if m % 2 == 0:
            result.append(-max_heap[0])

    print(len(result))
    cnt = 0
    for i in range(len(result)):
        if cnt == 10:
            print(); print(result[i], end = " "); cnt = 1
        else:
            print(result[i], end = " "); cnt += 1

결과