[BOJ] 1655 가운데를 말해요
우선순위 큐
이 문제는 하나의 입력이 주어질 때마다 여태까지의 값 들 중 중간값을 출력해야 하는 문제다.
그냥 시간복잡도 없이 생각을 해보자면, 하나씩 받아 리스트에 정렬하고 중간값을 출력하면 되는 문제이다. 그렇게 구현할 경우 O(N^2 x logN)의 어마어마한 시간복잡도를 얻게 된다.
이를 개선해보기 위해 N번만큼 정렬하는게 아닌, 매 인풋마다 정렬된 상태를 유지하면 된다. 그렇게 하려면 이분탐색을 사용하는 링크드리스트를 사용하거나, 우선순위큐를 사용하면 된다. 링크드리스트는 결국 중간값을 찾기 위해 처음부터 순차탐색을 해야하므로 시간복잡도가 복잡해진다.
그렇다면 우선순위큐로는 어떻게 해결할 수 있을까?
중간값을 뽑아내는 우선순위큐는 없다. 따라서 왼쪽과 오른쪽으로 나눠서 왼쪽에는 최대힙, 오른쪽에는 최소 힙을 구현하여 서로의 노드 개수 밸런싱을 해준다면 중간값을 계속 알 수 있지 않을까 하는 생각이다.
- 오른쪽 최소 힙의 탑 노드보다 작거나 같다면 왼쪽 최대 힙에 추가, 아니라면 오른쪽 최소힙에 추가
- 중간값을 얻기 위해 노드 개수를 밸런싱 한다.
- 왼쪽 힙의 크기가 2만큼 커지는 순간 하나를 Pop 해서 오른쪽 힙으로 이동시킨다.
- 반대로 오른쪽 힙의 크기가 1만큼 커지는 순간 하나를 pop 해서 왼쪽 힙으로 이동시킨다.
- 왼쪽 힙의 탑노드를 매번 출력한다.
이렇게 구현한다면 O(NlogN)으로 모든 중간값을 출력할 수 있게 된다.
구현
import sys
import heapq
if __name__ == "__main__":
N = int(sys.stdin.readline().rstrip())
leftMaxHeap, rightMinHeap = [10001], [10001]
leftCnt, rightCnt = 0, 0
for i in range(N):
num = int(sys.stdin.readline().rstrip())
# 새로운 값이 들어가야할 위치 선정
# 오른쪽보다 작거나 같으면 왼쪽에 삽입
if num <= rightMinHeap[0]:
heapq.heappush(leftMaxHeap, -num)
leftCnt += 1
# 아니라면 오른쪽에 삽입
else:
heapq.heappush(rightMinHeap, num)
rightCnt += 1
# 왼쪽,오른쪽 노드 개수 밸런싱
if leftCnt - rightCnt == 2:
tmp = -heapq.heappop(leftMaxHeap)
heapq.heappush(rightMinHeap, tmp)
leftCnt -= 1
rightCnt += 1
elif rightCnt > leftCnt:
tmp = heapq.heappop(rightMinHeap)
heapq.heappush(leftMaxHeap, -tmp)
rightCnt -= 1
leftCnt += 1
print(-leftMaxHeap[0])