[PS] BOJ 2696 / 중앙값 구하기
![[PS] BOJ 2696 / 중앙값 구하기](/content/images/size/w2000/2026/01/-------------1.png)
문제 링크: https://www.acmicpc.net/problem/2696
풀이
중앙값 찾기
어떤 수열에서 중앙값을 찾는 방법은 수 전체를 정렬하는 것입니다. 하지만, 수열에 원소를 계속 추가하면서 중앙값을 구해야 합니다. 매번 수를 추가할 때 마다 다시 정렬하는 방법은 분명 시간 초과를 받을 것이니, 더 빠른 방법을 찾아야 합니다.
매번 수를 추가할 때 마다 다시 정렬하는 대신, 우선순위 큐를 활용해 수가 추가될 때 정렬 순서에 맞는 위치에 들어가도록 구현할 수 있습니다. 우선순위 큐란, 가중치에 따라 특정 원소를 더 먼저 가져오는 형태의 큐입니다. 주로 힙(Heap) 자료구조를 기반으로 구현합니다.
하지만, 힙을 이용한 우선순위 큐의 구조를 생각해보면 가장 크거나 가장 작은 원소만 직접적으로 접근할 수 있습니다. 이런 큐에서 중앙값을 어떻게 구해야 할까요?
2개의 우선순위 큐 활용해 중앙값 찾기
우선순위 큐를 2개 사용해 중앙값을 즉시 찾을 수 있습니다.
전체 수열을 절반씩 나누어 각각의 힙(우선순위 큐)에 저장합니다. 이때, 작은 쪽 절반은 최대 힙으로, 큰 쪽 절반은 최소 힙으로 관리합니다.
from heapq import heappush, heappop
lower_half = [] # 작은 쪽 절반을 관리할 최대 힙
upper_half = [] # 큰 쪽 절반을 관리할 최소 힙두 개의 힙 선언하기
- 최대/최소 힙이란, 힙(우선순위 큐)에서 가장 먼저 꺼내질 원소가 최댓값/최솟값인 힙을 의미합니다.
- python의 heapq모듈은 최소 힙의 구현만 제공하므로, 최대 힙으로 사용하기 위해서 값을 음수로 바꿔 저장했습니다.
이후, 수열에 새로운 수를 추가하는 경우를 add_number 함수를 구현합니다. 이 함수는 새로운 수를 올바른 힙에 저장하고, 두 힙의 균형을 유지해줍니다.
- 두 힙의 크기는 항상 1 차이거나 같아야 합니다.
def add_number(num):
if len(lower_half) == 0 or num <= -lower_half[0]:
heappush(lower_half, -num) # heapq는 최소 힙으로 구현되어있으므로, 최대 힙으로 동작시키려면 음수로 저장해야 한다.
else:
heappush(upper_half, num)
# 두 힙의 균형을 맞춘다. (항상 두 힙의 크기 차이는 0 또는 1이어야 한다.)
if len(lower_half) > len(upper_half) + 1:
heappush(upper_half, -heappop(lower_half))
elif len(upper_half) > len(lower_half) + 1:
heappush(lower_half, -heappop(upper_half))두 개의 힙에 적절히 수 추가하고 관리하기
이제 두 개의 힙에서 중앙값을 구하는 함수 get_median을 구현합니다. 이 함수는 두 힙의 크기에 따라 중앙값을 찾습니다.
- 두 힙의 크기가 다르다면, 큰 쪽의 첫 원소가 중앙값입니다.
- 두 힙의 크기가 같다면, 전체 수열의 원소의 개수가 짝수 개인 상황이므로 중앙값이 2개입니다. 이 경우 문제 조건에 따라 적절히 구하면 됩니다.
- 평균값을 구하거나 둘 다 구하거나...
- 이 문제의 경우, 항상 수열의 원소의 개수가 홀수 개일 때에만 중앙값을 구하므로 별도로 고려하진 않았습니다.
def get_median():
if len(lower_half) > len(upper_half):
return -lower_half[0]
elif len(upper_half) > len(lower_half):
return upper_half[0]
else:
return -lower_half[0] # 짝수 개일 때는 작은 쪽의 최댓값을 반환한다.두 개의 힙에서 중앙값 구하기
전체 코드
from heapq import heappush, heappop
lower_half = []
upper_half = []
medians = [0] * 5000
def add_number(num):
if len(lower_half) == 0 or num <= -lower_half[0]:
heappush(lower_half, -num) # heapq는 최소 힙으로 구현되어있으므로, 최대 힙으로 동작시키려면 음수로 저장해야 한다.
else:
heappush(upper_half, num)
# 두 힙의 균형을 맞춘다. (항상 두 힙의 크기 차이는 0 또는 1이어야 한다.)
if len(lower_half) > len(upper_half) + 1:
heappush(upper_half, -heappop(lower_half))
elif len(upper_half) > len(lower_half) + 1:
heappush(lower_half, -heappop(upper_half))
def get_median():
if len(lower_half) > len(upper_half):
return -lower_half[0]
elif len(upper_half) > len(lower_half):
return upper_half[0]
else:
return -lower_half[0] # 짝수 개일 때는 작은 쪽의 최댓값을 반환한다.
def solve():
lower_half.clear()
upper_half.clear()
N = int(input())
res_size = (N + 1) // 2
start = 0
while N > 0:
stream = map(int, input().split())
for i, v in enumerate(stream, start=start):
add_number(v)
if i % 2 == 0: # 인덱스가 짝수일 때 (홀수 번째)
medians[i // 2] = get_median()
start += 10
N -= 10
start = 0
print(res_size)
while start < res_size:
end = min(start + 10, res_size)
print(' '.join(map(str, medians[start:end])))
start += 10
input = open(0).readline
for _ in range(int(input())):
solve()solution.py