[BOJ] 1655 가운데를 말해요

우선순위 큐


이 문제는 하나의 입력이 주어질 때마다 여태까지의 값 들 중 중간값을 출력해야 하는 문제다.

그냥 시간복잡도 없이 생각을 해보자면, 하나씩 받아 리스트에 정렬하고 중간값을 출력하면 되는 문제이다. 그렇게 구현할 경우 O(N^2 x logN)의 어마어마한 시간복잡도를 얻게 된다.

이를 개선해보기 위해 N번만큼 정렬하는게 아닌, 매 인풋마다 정렬된 상태를 유지하면 된다. 그렇게 하려면 이분탐색을 사용하는 링크드리스트를 사용하거나, 우선순위큐를 사용하면 된다. 링크드리스트는 결국 중간값을 찾기 위해 처음부터 순차탐색을 해야하므로 시간복잡도가 복잡해진다.

그렇다면 우선순위큐로는 어떻게 해결할 수 있을까?

중간값을 뽑아내는 우선순위큐는 없다. 따라서 왼쪽과 오른쪽으로 나눠서 왼쪽에는 최대힙, 오른쪽에는 최소 힙을 구현하여 서로의 노드 개수 밸런싱을 해준다면 중간값을 계속 알 수 있지 않을까 하는 생각이다.

  1. 오른쪽 최소 힙의 탑 노드보다 작거나 같다면 왼쪽 최대 힙에 추가, 아니라면 오른쪽 최소힙에 추가
  2. 중간값을 얻기 위해 노드 개수를 밸런싱 한다.
    1. 왼쪽 힙의 크기가 2만큼 커지는 순간 하나를 Pop 해서 오른쪽 힙으로 이동시킨다.
    2. 반대로 오른쪽 힙의 크기가 1만큼 커지는 순간 하나를 pop 해서 왼쪽 힙으로 이동시킨다.
  3. 왼쪽 힙의 탑노드를 매번 출력한다.

이렇게 구현한다면 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])