[ BOJ 1655 ] 가운데를 말해요(Python)

문제

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

아주 그냥 요상한 게임은 다 한다.
정수가 하나씩 추가될 때마다, 지금까지 추가된 수를 정렬한 뒤 가운데에 있는 수를 출력하면 된다.

문제 풀이

[오답]
힙으로 입력을 받아서 자동으로 정렬하고, 가운데에 있는 요소를 출력하면 되겠다고 생각했다.

아래와 같은 코드로 접근했다가 시간 초과가 발생했다.

for i in range(n):
    num = int(input())
    heapq.heappush(h, num)
    if len(h) < 3:
        print("a:",(h[0]))
    else:
    	# 중앙값 = 배열의 절반만큼 pop
        k = ((len(h) + 1) // 2)
        hc = h.copy()
        for _ in range(k-1):
            heapq.heappop(hc)
        print("a:",heapq.heappop(hc))

입력 개수는 100,000개로 그렇게 많지 않지만, 시간 제한이 0.1초다.
최소힙일 때 가장 작은 값에 접근하는 것 자체는 O(1)이지만, 가장 작은 요소를 pop 하거나, 다른 요소를 push 할 때 내부에서 정렬을 하는데에 걸리는 시간은 O(n log n) 만큼 소요된다.

[정답]
풀이를 찾아보니 힙을 두 개 써서 n을 분배할 수 있었다.

왼쪽 힙은 최대 힙으로 정렬하고, 오른쪽 힙은 최소 힙으로 정렬하면
왼쪽 힙의 첫번째 요소는 항상 중앙값이 된다.

• 이때 왼쪽 힙의 첫번째 요소는 최대 힙에서 가장 큰 값이다.
• 마찬가지로 오른쪽 힙의 첫번째 요소는 최소 힙에서 가장 작은 값

1. 왼쪽 힙과 오른쪽 힙의 길이가 같으면 (요소 * -1) 을 왼쪽 힙에 삽입한다.
   1-1. 그렇지 않으면 오른쪽 힙에 삽입한다.
2. 왼쪽 힙과 오른쪽 힙에 요소가 존재하고, 왼쪽 힙의 (첫번째 요소* -1) 가 오른쪽 첫번째 요소보다 클 때
   2-1. 왼쪽 힙의 첫번째 요소와 오른쪽 힙의 첫번째 요소를 바꿔준다. ( -1을 곱해준 뒤 바꿔준다. )
3. 왼쪽 힙의 (첫번째 요소 * -1)를 출력한다.

left = 최대 힙(내림차순), right = 최소 힙(오름차순)

ex) [1, 5, 2, 10, -99, 7, 5]
num = 1 👉🏻 left = [-1] / right = []
num = 5 👉🏻 left = [-1], right = [5]
num = 2 👉🏻 left = [-2,-1], right = [5]
num = 10 👉🏻 left = [-2,-1], right = [5,10]
num = -99 👉🏻 left = [-2,-1,99], right = [5,10]
num = 7 👉🏻 left = [-2,-1,99], right = [5,7,10]
num = 5 👉🏻 left = [-5,-2,-1,99], right = [5,7,10]

코드

import sys
import heapq
input = sys.stdin.readline

n = int(input())
max_h, min_h = [], []

for i in range(n):
    num = int(input())
    if len(max_h) == len(min_h):
        heapq.heappush(max_h, -num)
    else:
        heapq.heappush(min_h, num)

    if len(max_h) >= 1 and len(min_h) >= 1 and max_h[0] * -1 > min_h[0]:
        max_value = heapq.heappop(max_h) * -1
        min_value = heapq.heappop(min_h)
        
        heapq.heappush(max_h, min_value * -1)
        heapq.heappush(min_h, max_value)

    print(max_h[0] * -1)