BOJ 1655 가운데를 말해요 _ 파이썬

📃 문제

BOJ _ 1655 가운데를 말해요

알고리즘 - 자료구조, 우선순위 큐

문제

백준이가 N개의 숫자를 순서대로 외치면 동생은 매번 지금까지 외친 수의 중간값을 말한다.
외친 수의 갯수가 짝수개라면 중간에 있는 두 수 중 작은 수를 말한다.

입력

7
1
5
2
10
-99
7
5

출력

1
1
2
2
2
2
5

조건

파이썬 기준 시간 제한이 0.6초
-> 매번 배열을 정렬하는 등의 방법은 불가능!

💡 풀이 방법

힙큐(우선순위 큐)를 사용하고 두 개의 힙큐가 필요하다!

중앙값을 기준으로 그보다 작거나 같은 수가 들어갈 작은수 힙과 그보다 큰 수가 들어갈 큰 힙이 필요하다.
그리고 중앙값은 "작은 수 힙 중 가장 큰 값" 이다!

1. 매번 백준이가 외치는 수에 대해 크기 비교를 통해 맞는 힙에 넣어준다.
2. 길이를 계산하여 대칭을 맞춰준다.
3. 이 떄 중앙값을 갱신하는 작업이 가능하다!! 길이는, 작은수 힙 길이 - 큰수 힙 길이 == 0 or 1이어야 한다.
   -> 왜냐하면, "중앙값은 작은수힙 중 가장 큰값"으로 두어야하기 때문이다.

두 힙의 길이가 같다면 (ex 1 3 / 5 6) 작은힙 중 가장 큰 값이 중앙값임이 만족된다.
작은 힙의 길이가 하나 크담녀 (ex 1 3 4 / 5 6) 마찬 가지로 작은 힙 중 가장 큰 값이 중앙값임이 만족된다.

코드

import sys
input = sys.stdin.readline
import heapq

hq = []
n = int(input())

x= int(input())     # 처음 입력값은 작은수힙에 넣고 중앙값이다.
smallerheap =[-x]    # 작은수힙은 최대힙으로 유지하기 위해 음수값으로 저장한다.
biggerheap = []
mid = x
print(mid)

for _ in range(n-1):
    x = int(input())
    if x>= mid:     # 중앙값보다 큰 입력이 주어진 경우 큰수힙에 넣는다.
        heapq.heappush(biggerheap,x)
        # 이때, 두힙의 길이가 같아지거나 큰수힙이 하나 더 길어지는 경우만 존재한다.
        if len(smallerheap) == len(biggerheap):
            mid = -smallerheap[0]   # 길이 같은 경우, 최소힙에서의 가장 큰값이 중앙값. 음수로 힙에 저장했었기 때문에 -를 붙여서 반환
        elif len(biggerheap) - len(smallerheap) == 1:
            mid = heapq.heappop(biggerheap) # 큰수힙이 하나 더 길어진 경우, 큰수힙에서 가장 작은 값이 중앙값이고, 길이를 맞추기 위해 작은수힙에 음수 붙여 저장
            heapq.heappush(smallerheap,-mid)
    else: # x < mid # 중앙값보다 작은 입력이 주어진 경우 작은수힙에 넣는다.
        heapq.heappush(smallerheap,-x)  #작은수 힙은 최대힙으로 저장하기 위해 음수값으로 저장한다.
        # 이때, 작은수힙이 하나 더 길거나, 두개 더 긴 경우만 존재한다.
        if len(smallerheap) - len(biggerheap) >=2:
            # 두개 긴 경우 작은수힙의 가장 큰값은 큰수힙으로 옮기고 그다음 가장 큰수가 mid가 된다.
            heapq.heappush(biggerheap, -1 * heapq.heappop(smallerheap))
            mid = -smallerheap[0]
        # 하나 더 긴 경우는, 이전 단계에서 서로 길이가 같았다는 것이고 그렇다면 이미 최소힙의 가장큰값을 중앙값으로 둔 상태이니 mid에 갱신은 없다.
    print(mid)

... 추가 반례

'''
1
2
3
3
3
5

>>
1
1
2
2
3
3


5
3
2
4
1
5

>>
3
2
3
2
3
'''

🔍 정리 & 배운 것

시간 복잡도 효율적인 중앙값 찾기는 두 개의 힙큐 ! 기억하자 !!