[Python] 프로그래머스 / 이중우선순위큐

문제 링크

heapq의 함정

자료구조 시간에 다들 Max Heap과 Min Heap 구조에 대해 배웠을 것이다. Tree 구조를 이용해서 리스트에서 최소값이나 최대값을 찾는 시간을 O(logN)으로 줄인 획기적인 기법이다. 이를 파이썬에서 쉽게 사용할 수 있도록 지원해주는 라이브러리가 heapq이다. (Min Heap만 가능) 그런데, 코테에서 아무 생각 없이 heapq를 자주 사용하다보면, heapq를 그저 시간효율성이 좋은 sort()로 착각하기 쉽다는 것이다!

 

여기 예시가 있다.

import heapq
a = [21, 20, 8, 19, 1, 14, 6, 8, 30, 26]
heapq.heapify(a)
a
#[1, 8, 6, 19, 20, 14, 8, 21, 30, 26]

 

위에서 보면, a[0]가 최소값이 것이 언제나 보장되지만, 그 이후에는 꼭 크기순으로 배열되어있지는 않다. heap 자료구조형은, 조부모/부모/자식/손자 노드 등 직계 노드와의 부등호만 성립할 뿐, 방계로는 부등호가 성립하지 않기 때문이다.

 

그러므로, 어떤 리스트를 heapq로 관리할 때, 최대값을 찾으려면 nlargest()를 사용해야 한다.

heapq.nlargest(1, a)
#[30]

 

여기서 nlargest의 첫 번째 파라미터는, 큰 순으로 몇 등까지 찾고싶은가를 정하는 것이다.

 

정답 코드 

import heapq

def solution(operations):
    queue = []
    
    for operation in operations:
        order, number = operation.split()
        number = int(number)
        
        if order == "I":
            heapq.heappush(queue, number)
        elif order == "D" and number == -1:
            if queue:
                heapq.heappop(queue)
        elif order == "D" and number == 1:
            if queue:
                queue.remove(heapq.nlargest(1, queue)[0])
    
    if not queue:
        return [0, 0]
    else:
        max_value = heapq.nlargest(1, queue)[0]
        min_value = queue[0]
        return [max_value, min_value]