[Python] 백준/BOJ 1655번: 가운데를 말해요 (Gold 2)

💻 Problem

문제 보러 가기

 

백준이는 동생에게 "가운데를 말해요" 게임을 가르쳐주고 있다. 백준이가 정수를 하나씩 외칠 때마다 동생은 지금까지 백준이가 말한 수 중에서 중간값을 말해야 한다. 만약, 그동안 백준이가 외친 수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다.

예를 들어 백준이가 동생에게 1, 5, 2, 10, -99, 7, 5를 순서대로 외쳤다고 하면, 동생은 1, 1, 2, 2, 2, 2, 5를 차례대로 말해야 한다. 백준이가 외치는 수가 주어졌을 때, 동생이 말해야 하는 수를 구하는 프로그램을 작성하시오.

 

💡 Approach

백준이가 숫자를 하나씩 말할 때마다 지금까지 나온 수들 중에 중간값을 말해야 한다.

지금까지 나온 숫자의 개수가 홀수면 가운데 값, 짝수면 가운데 두 개 중 작은 값이다.

순서	누적된 수	정렬 결과	중간값
1	[1]	[1]	1
2	[1, 5]	[1, 5]	1
3	[1, 5, 2]	[1, 2, 5]	2
4	[1, 5, 2, 10]	[1, 2, 5, 10]	2
5	[1, 5, 2, 10, -99]	[-99, 1, 2, 5, 10]	2
6	[1, 5, 2, 10, -99, 7]	[-99, 1, 2, 5, 7, 10]	2
7	[1, 5, 2, 10, -99, 7, 5]	[-99, 1, 2, 5, 5, 7, 10]	5

 

이중 힙(최대 힙 + 최소 힙)을 사용해서 중간값을 관리했다.

 

• 최대 힙 (left): 중간값 이하의 숫자 저장
• 최소 힙 (right): 중간값 초과 숫자 저장

이중 힙의 밸런스를 조정하여 항상 left 쪽의 top을 중간값으로 유지한다.

 

✏️ Solution

import sys
import heapq
input = sys.stdin.readline

N = int(input())
left = [] # 최대 힙 (중간값 포함)
right = [] # 최소 힙

for i in range(N):
    num = int(input())

    # 숫자가 중간값 이하면 left(최대 힙)에 삽입한다
    if not left or num <= -left[0]:
        heapq.heappush(left, -num)
    # 숫자가 중간값보다 크면 right(최소 힙)에 삽입한다
    else:
        heapq.heappush(right, num)

    # 힙 밸런스 조정
    # mid를 기준으로 left 개수보다 right 개수가 더 많으면
    # right 값을 하나 빼서 left로 이동한다
    if len(left) < len(right):
        heapq.heappush(left, -heapq.heappop(right))
    # right 개수 + 1 보다 left 개수가 더 많으면
    # left 값을 하나 빼서 right로 이동한다
    elif len(left) > len(right) + 1:
        heapq.heappush(right, -heapq.heappop(left))

    print(-left[0])

우선순위 큐의 시간 복잡도는 logN이고 N번 반복하니 전체 시간 복잡도는 O(NlogN)이다.

N의 최댓값은 100,000이다.

10^6은 2^20으로 치환 가능하므로 시간 복잡도는 2 * 10^7이다.