1655번: 가운데를 말해요 (우선순위 큐)

1655번: 가운데를 말해요

문제자체가 매우 간단해보여서 리스트를 이용해 정렬한 후 중간값을 출력했는데 시간초과가 떴다. 정렬을 하는데 O(nlogn) 걸리고 이를 N번 반복한다고 하면 대충봐도 O(N^2)은 가뿐히 넘긴다. 여기서 N이 1보다 크거나 같고, 100,000보다 작거나 같으므로 시간초과가 뜨는 것은 당연하다. 따라서 매 입력이 들어올 때마다 O(nlogn) 보다 작은 시간복잡도를 가지는 방법으로 중간값을 찾아야 한다. 전혀 시간복잡도를 줄일 방법이 떠오르지 않아 이 문제의 알고리즘 분류를 봤는데 '우선순위 큐'가 있어서 자료구조 중 하나인 우선순위 큐를 공부하고 이를 최대 힙, 최소 힙을 이용해 중간값을 찾는 알고리즘을 떠올렸다. 

 

문제공부내용 (엔지니어대한민국, 이코테 설명을 참고함):

• 우선순위 큐(Priority Queue) 는 들어간 순서에 상관없이 우선순위가 높은 데이터가 먼저 나오는 것을 의미한다.
• 주로 완전이진트리를 가지는 heap 자료구조를 이용해 구현을 하며 루트 노트에 우선순위가 높은 것이 위치한다.
• 최대 힙(Max Heap)은 완전이진트리이며 루트 노드로 올라가며 저장된 값이 커진다 
• 최소 힙(Min Heap)은 완전이진트리이며, 루트 노드로 올라가면 저장된 값이 작아진다.
• 파이썬에서는 우선순위큐를 heapq 모듈을 이용해 구현한다. 파이썬은 기본적으로 최소힙으로 설정이 되어있다. 따라서 최대힙을 구현할 때는 -를 붙여서 이용한다 (밑에 문제 풀이에 사용한 코드를 참고)
• 데이터의 개수가 N개일 때, Heap으로 우선순위 큐를 구현했을 때 삽입시간은 O(logN), 삭제시간은 O(logN)이다. 
• 단순히 N개 데이트를 힙에 넣어다가 모두 꺼내는 작업은 정렬과 동일 (힙 정렬)

문제 접근 및 해설 : 

우선순위 큐를 배우고 구현하는 방법까지 찾아봤는데 그럼 어떻게 중간값을 탐색할 수 있을지 고민해봤고 중간값을 결정짓는 조건을 생각해보면서 아래와 같은 생각을 했다. 

1. 자연수들로 이루어진 두 집합 A,B가 있다고 생각하자. A의 모든 원소가 B의 모든 원소보다 작을 때, n(A) = n(B)+1 이거나 n(A) = n(B) 이면 A의 최댓값이 문제에서 구하고자 하는 중간값이다. (짝수개일때는 크기가 작은 값이 중간값으로 설정한다고 문제에서 제시됨) 
2. 자연수들로 이루어진 집합 A,B가 있다고 생각해보자. A에서의 최댓값을 m, B에서의 최솟값을 n이라고 할때, m < n 이라면 A의 원소들은 B의 원소들보다 모두 작은 값을 갖고 있다고 생각할 수 있다. 이때 1을 이용하여 중간값을 찾아낼 수 있다. 
3. 또한 현재 A의 모든 원소가 B의 모든 원소보다 작은 상태를 생각해보자. A에다가 어떤 원소를 넣었을 때, A에서의 최댓값이 m, B에서의 최솟값을 n이라고 하면 m < n인 경우에는 2번과 동일하고, m > n인 경우에는 A에 추가한 원소가 B의 최솟값보다 큰 경우이므로 A에 추가한 원소와 B의 최소값을 서로 교환하면 다시 A의 모든 원소가 B의 모든 원소보다 작은 상태를 만들 수 있다. 따라서 1을 이용하여 중간값을 찾아낼 수 있다.

위 아이디어를 보면서 어떤 생각이 떠오르지 않는가?

• 바로 A는 최대 힙, B는 최소 힙을 이용하여 구현하면 될 듯.
• 그리고 n(A)를 n(B) 혹은 n(B) + 1을 계속 유지하여야 하는데 입력받은 수를 A, B 에 번갈아 넣어주면 될 듯. 
• 매 입력이 들어올 때마다 위에서 '3' 을 반복하기만 하면 계속 A의 모든 원소가 B보다 작게 유지할 수 있다. 

작성 코드(Python3):

 

import sys
import heapq
input = sys.stdin.readline

n = int(input())
A = []
B = []
for i in range(n):
    x = int(input())
    if len(A)==len(B):
        heapq.heappush(A,-x) # 최대힙을 구현하기 위해 -를 붙임
    else:
        heapq.heappush(B,x)
    if not B:
        print(-A[0])
        continue
    if -A[0] > B[0]:
        n = -heapq.heappop(A)
        m = heapq.heappop(B)
        heapq.heappush(A,-m)
        heapq.heappush(B,n)
    print(-A[0])