[백준] 1655 가운데를 말해요 (python) - 중앙값 힙(Median Heap)

 

1️⃣ 문제 링크

 

https://www.acmicpc.net/problem/1655

 

2️⃣ 기억해야 했던 개념들

 

- 매번 median 값을 어떻게 빠르게 구할 수 있을까? (python 기준 0.6초)

- 정렬을 매번 하면 O(N² log N)이므로 10⁵ 데이터면 시간초과. 

- 정렬을 하지 않고, 중앙값을 찾을 수 있는 방법을 생각 -> 중앙값 힙(Median Heap) 개념을 이용한다. 

 

3️⃣ 풀이 과정

두 그룹으로 나누기

• 중앙값보다 작거나 같은 값 → 왼쪽 그룹 (MAX HEAP)
• 중앙값보다 큰 값 → 오른쪽 그룹 (MIN HEAP)

MAX_HEAP의 크기 = MIN_HEAP 크기 + 1 로 유지한다.

• 들어오는 값이 현재 중앙값이 작거나 같으면 MAX HEAP에 넣고 아니면 MIN HEAP에 넣는다.
• 만약 크기 차가 1보다 커지면 heap에서 pop해서 다른쪽 heap에 push한다.

중앙값 구하기

• MAX HEAP의 0번 값을 읽으면 항상 중앙값이다.

4️⃣ 코드

import heapq
import sys
input = sys.stdin.readline

N = int(input())

max_que = []
min_que = []

for i in range(N):
    num = int(input())
    if (len(max_que) == 0 or num <= -max_que[0]):
        heapq.heappush(max_que, -num)
    else:
        heapq.heappush(min_que, num)

    if len(max_que) > len(min_que) + 1:
        heapq.heappush(min_que, -heapq.heappop(max_que))
    elif len(min_que) > len(max_que):
        heapq.heappush(max_que, -heapq.heappop(min_que))
    
    print(-max_que[0])

 

 

유명한 알고리즘이라는데 처음 알았다..