[BOJ] 2696 - 중앙값 구하기 (Python)

 

 

아이디어

 

최소힙과 최대힙의 두 가지 우선순위 큐를 사용하여 중앙값을 찾는다.

 

 

풀이

 

최대힙 left와 최소힙 right를 빈 리스트로 생성하고, 처음 입력 받는 값을 중앙값 mid로 저장한다.

이후 중앙값과 비교해 작거나 같은 값은 최대힙 left에, 큰 값은 최소힙 right에 push한다.

두 힙의 크기가 2 이상 차이나는 경우 현재 중앙값은 크기가 작은 힙에 push하고,

크기가 큰 힙에서 루트 노드를 pop하여 중앙값을 새롭게 바꾸고 힙의 균형을 맞춘다.

홀수 번째 입력 때마다 중앙값을 결과값에 숫자를 문자열로 추가하고, 조건에 따라 개행 문자를 추가한다.

 

import sys, heapq

input = sys.stdin.readline

T = int(input())
for t in range(T):
    M = int(input())
    left = []
    right = []
    for i in range(M//10+1):
        nums = list(map(int, input().split()))
        for j in range(len(nums)):
            if not i and not j:
                mid = nums[0]
                res = f'{mid}'
                cnt = 1
                continue
            if nums[j] <= mid:
                heapq.heappush(left, -nums[j])
            else:
                heapq.heappush(right, nums[j])
            if len(left) > len(right)+1:
                heapq.heappush(right, mid)
                mid = -heapq.heappop(left)
            elif len(left)+1 < len(right):
                heapq.heappush(left, -mid)
                mid = heapq.heappop(right)
            if not j%2:
                res += f' {mid}' if cnt%10 else f'\n{mid}'
                cnt += 1
    print(cnt, res, sep='\n')

 

 

결과값을 문자열이 아닌 리스트로 받고 for문에서 슬라이싱하여 언패킹할 수도 있다.

큰 차이는 아니지만 메모리가 150KB 정도 더 사용되기에 리스트가 아닌 문자열로 받도록 했다.

res = [mid]
...
if not j%2:
    res.append(mid)
...
print(len(res))
for i in range(len(res)//10+1):
    print(*res[i*10:(i+1)*10])