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

-10,000 ~ 10,000 사이의 정수 N개가 입력으로 들어오면, 숫자 입력이 차례로 들어올 때마다 중앙값을 출력하는 문제다.

ex) 
input : 3 6 8 4 5 6 2
output : 3 3 6 4 5 5 5

입력이 짝수개가 들어왔을 때는 중간의 두 수중 작은 수를 출력한다.

시간 복잡도: O(NlogN)

---

풀이 방법

출력해야 하는 값은 중앙값이다. 제일 간단히 중앙값을 찾는 방법은, 입력이 들어올 때마다 들어온 입력값들을 정렬한 후에 가운데 값을 출력하는 것이지만 그러면 매 입력마다 O(NlogN)의 시간이 걸리고 총 O(N^2 * logN)의 시간이 걸려서 시간 초과가 발생한다.

Heap 자료구조를 이용하면 삽입, 삭제를 O(logN) 시간에 할 수 있으며 구현 방법에 따라 MinHeap, MaxHeap을 만들어서 각각 Heap 자료구조 안의 원소 중 가장 작은 값, 가장 큰 값을 가져올 수 있다.

이때 Heap을 이용한 우선순위 큐(Priority Queue) 2개를 만들어(MaxHeap, MinHeap) 각 우선순위 큐의 사이즈를 같거나 1개까지만 차이 나도록 하면 큐의 첫 번째 요소를 통해 중앙값을 찾을 수 있다.
※ Heap은 우선순위 큐를 구현하기 위해 사용되는 자료구조다. 이 포스팅에서는 두 용어를 혼합해 사용한다.

예시
입력이 3 6 8 4 5 6 2로 들어온다고 가정해보자

입력이 들어올 때 값들을 2개의 우선순위 큐에 넣을 것이다. 이때 작은 값 절반은 Max Heap에 넣을 것이고 큰 값 절반은 MinHeap에 넣을 것이다. 이때 입력으로 들어온 숫자 개수가 홀수일 때는 Max Heap의 길이가 Min Heap의 길이보다 1 크다.
이렇게 되면 항상 Max Heap의 첫 번째 원소가 중앙값이 된다. 또한 직관성을 위해 Max Heap은 위치를 거꾸로 표시해 놨으며, Heap의 자료구조 특성상 첫 번째 원소를 제외하고는 위치에 따른 원소 간 대소 관계가 일정하지 않다. 

첫번째 입력이 들어왔을 때

첫 번째 입력이 들어왔을 때는 Max Heap에 원소를 삽입한다. 총입력이 들어온 개수가 홀수이므로 Max Heap의 길이가 Min Heap의 길이보다 1 크다.

두번째 입력이 들어왔을 때

두 번째 입력이 들어올 때는 Max Heap의 크기가 Min Heap보다 1 큰 상황이다. 이때는 다음 과정을 거친다.

1. Max Heap에서 원소를 pop 한다. (Max Heap의 가장 큰 원소가 빠진다.)
2. 입력으로 들어온 값과 pop 한 원소를 비교해서 더 작은 값을 Max Heap에 삽입한다.
3. 나머지 값을 Min Heap에 삽입한다.

이 과정을 거치면 Max Heap과 Min Heap의 크기가 같아지며 Max Heap의 모든 원소가 Min Heap의 모든 원소보다 작거나 같은 조건이 유지된다.

세번째 입력이 들어왔을 때

세 번째 입력이 들어올 때는 Max Heap의 크기와 Min Heap의 크기가 같은 상황이다. 이때는 다음의 과정을 거친다.

1. 입력값과 Min Heap의 첫 번째 요소를 비교한다. 입력값이 더 크다면, Min Heap에 입력값을 push 하고 Min Heap에서 원소를 pop 한다. 크지 않다면 바로 다음 단계로 넘어간다.
2. 입력값이 Min Heap의 첫 번째 요소보다 컸다면 pop 한 원소를 Max Heap에 삽입한다. 크지 않았다면 입력값을 Max Heap에 삽입한다.

이 과정을 거치면 Max Heap의 크기가 Min Heap보다 1 커지게 된다.

네번째 입력이 들어왔을 때

네 번째 입력이 들어왔을 때는 두 번째 입력이 들어왔을 때와 같은 상황이다. 같은 과정을 통해 Heap에 원소를 삽입한다.

다섯번째 입력이 들어왔을 때

다섯 번째 입력이 들어왔을 때는 세 번째 입력이 들어왔을 때와 같은 상황이다. 같은 과정을 통해 Heap에 원소를 삽입한다.

 

이후 과정 역시 다음과 같다.

여섯번째 입력이 들어왔을 때

일곱번째 입력이 들어왔을 때

요약하면 홀수번째 입력이 들어올 때는 Max Heap의 크기가 Min Heap의 크기보다 1 커지며, 짝수번째 입력이 들어올 때는 두 Heap의 크기가 같아진다. 또한 입력이 들어왔을 때 실행되는 로직이 순서에 따라 2개로 나뉜다.
입력이 들어올 때마다 Heap의 삽입, 삭제가 1~3번 발생하므로 O(logN)의 시간으로 처리가 가능하다. 입력이 N개 들어오므로 총 시간 복잡도는 O(NlogN)이다.

---

코드

import sys
import heapq
 
input = sys.stdin.readline
 
N = int(input())
 
minpq = []
maxpq = []
 
for i in range(N):
    num = int(input())
    if len(minpq) == len(maxpq):
        if len(minpq) > 0 and minpq[0] < num:
            heapq.heappush(minpq, num)
            num = heapq.heappop(minpq)
        heapq.heappush(maxpq, -num)
    elif len(minpq) < len(maxpq):
        temp = -heapq.heappop(maxpq)
        if num <= temp:
            heapq.heappush(maxpq, -num)
            heapq.heappush(minpq, temp)
        else:
            heapq.heappush(minpq, num)
            heapq.heappush(maxpq, -temp)
 
    print(-maxpq[0])

 

주석 추가

import sys
import heapq
 
input = sys.stdin.readline  # input 효율화
 
N = int(input())  # 입력 개수
 
minpq = []  # 최소 힙
maxpq = []  # 최대 힙
 
for i in range(N):
    num = int(input())  # 입력값
    if len(minpq) == len(maxpq):  # 두 Priority Queue의 크기가 같다면(홀수번째 입력)
        if len(minpq) > 0 and minpq[0] < num:  # MinHeap의 첫번째 요소보다 입력값이 크다면
            heapq.heappush(minpq, num)  # 입력값을 MinHeap에 삽입
            num = heapq.heappop(minpq)  # MinHeap에서 원소를 pop (후에 MaxHeap에 삽입하기 위해 num 변수의 값을 대체)
        heapq.heappush(maxpq, -num)  # MaxHeap에 값 삽입 (heapq모듈은 MinHeap으로 작동되기 때문에 입력값이 음수 곱해서 삽입)
    elif len(minpq) < len(maxpq):  # MaxHeap의 크기가 더 크다면(짝수번째 입력)
        temp = -heapq.heappop(maxpq)  # MaxHeap에서 값 pop
        if num <= temp:  # 입력값이 pop한 값보다 작거나 같다면
            heapq.heappush(maxpq, -num)  # MaxHeap에 입력값 삽입
            heapq.heappush(minpq, temp)  # MinHeap에 pop한 값 삽입
        else:
            heapq.heappush(minpq, num)  # MinHeap에 입력값 삽입
            heapq.heappush(maxpq, -temp)  # MaxHeap에 pop한 값 삽입
 
    print(-maxpq[0]) #중앙값 출력 (중앙값은 항상 MaxHeap의 첫번째 요소)

---

채점 결과

전에 시도했다가 풀지 못했던 문제였다. 이번에 시도할 때도 조건을 몇 개 놓쳐서 4번 만에 성공했다.