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

🧑🏻‍💻 문제링크

문제풀이

파이썬의 heapq 모듈은 기본적으로 Min Heap 구조를 갖는다. 그렇기 때문에 heappop()을 하게되면 가장 heap에서 최솟값을 빼게 된다.

처음에 생각했던 방법은 max_heap, min_heap 두 개의 heap을 만들어서 각각 최댓값과 최솟값을 구하는 방식이었다. 그렇기 위해서는 최대 힙을 구해야 하는데 최대 힙을 구하는 방식은 다음과 같다.

import heapq

data = [2, 3, 1, 4, 9, 7]
heap = []

for num in nums:
  heapq.heappush(heap, (-data, data))  # (우선 순위, 값)

while heap:
  print(heapq.heappop(heap)[1])  # index 1

이런 방식으로 튜플(Tuple)을 이용하여 힙에 추가하거나 삭제하면 된다. 하지만 이렇게 하면 코드가 복잡해지기 때문에 heapq 모듈을 좀 더 구글링 해봤다.

몇 분 간의 구글링해서 문제를 간단하게 해결할 수 있는 방법을 찾을 수 있었다.🤣

heapq 모듈에 있는 nlargest()와 nsmallest() 함수를 사용해서 최대 또는 최소값을 찾을 수 있다.

heapq.nlargest(n, iterable, key=None)
heapq.nsmallest(n, iterable, key=None)

하지만 굳이 최솟값을 찾는데 nsmallest() 까지 사용할 필요는 없고 최댓값을 찾은 후 다시 heapify()로 힙에 넣어주면 된다.

코드

import heapq

def solution(operations):
    answer = []
    heap = []   # 힙 생성

    for data in operations:
        # 연산이 "I"일 경우 공백 뒤의 숫자를 heap에 넣음
        if data[0] == "I":
            heapq.heappush(heap, int(data[2:]))
        # 연산이 "D"일 경우
        else:
            if len(heap) == 0:
                pass
            # 공백 뒤가 "-"일 경우 최소힙을 제거
            elif data[2] == "-":
                heapq.heappop(heap)
            # 공백 뒤가 아무것도 아닌 경우
            else:
                # 힙에서 가장 큰 수를 제외하고 다시 힙에 넣음
                heap = heapq.nlargest(len(heap), heap)[1:]
                heapq.heapify(heap)

    if heap:
        answer.append(max(heap))
        answer.append(min(heap))
    else:
        answer.append(0)
        answer.append(0)
    
    return answer

결과