[Programmers] 힙(Heap) - 이중우선순위큐 (Python)

 

코딩테스트 연습 - 이중우선순위큐

 

오랜만에 기초 문제 다시 풀기!

heap 자료구조를 사용한 문제풀이는 정말 오랜만이라서 heapq 라이브러리의 간단한 사용법을 정리할겸 써본다.

 

파이썬에서는 내장 라이브러리로 heapq라는 기능을 제공하기 때문에 heap 문제를 매우 간단하게 풀 수 있다.

우선 heap을 사용해서 풀 때 유용한 유형은 "최댓값/최솟값을 반복적으로 리턴해야할 때"이다!

 

사용법은 간단하다.

 

import heapq

heap = [] # 그냥 리스트로 힙을 만들어준다
heapq.heappush(heap, 1)
heapq.heappop(heap)​

 

그냥 리스트로 힙을 만들어준 후 heapq라이브러리를 사용하여 값을 넣고, 빼주면 된다는 것!

 

만약 풀고자하는 문제가 

1. 리스트에서 최솟값을 항상 뽑아야 한다.

위 예시처럼 heappush로 값을 넣고, heappop으로 그냥 빼주면 된다.

 

 

2. 리스트에서 최댓값을 항상 뽑아야 한다.

heapq.heappush(heap, -num)

heap에 숫자를 넣을 때 항상 음수로 넣어주면 최댓값을 얻을 수 있다!

 

 

3. 어쩔때는 최댓값, 어쩔때는 최솟값을 뽑아야 한다면..?

이게 바로 프로그래머스에서 제공하는 이중우선순위큐 유형의 문제이다!

이런 경우에는 2번처럼 음수로 숫자를 저장해서는 안된다! (최소값을 빼는 경우도 있으니까..)

대신 heapq 라이브러리의 nlargest라는 기능을 사용하면 된다.

nlargest란 heap에서 숫자가 큰 순서대로 n개를 뽑아주는 기능이다.

이때 큰 순서대로 정렬하여 리턴해준다!

 

heap = heapq.nlargest(len(heap), heap)[1:]

 

따라서 heap의 전체 길이만큼 nlargest를 해주면 전체 heap이 큰 순서대로 정렬이 된다.

여기서 [1:]로 슬라이싱을 해주면 제일 큰 숫자를 제외하고 나머지 숫자만이 heap에 남게 된다.

 

nlargest를 해준 이후에는 heapify를 통해 다시 heap의 구조로 정렬을 꼭 해주어야 한다!

heapq.heapify(heap)

 

 

<프로그래머스 - 이중우선순위큐 풀이>

 

import heapq

def solution(operations):
    heap = []
    for operation in operations :
        op0, op1 = operation.split()
        if op0 == "I" :
            heapq.heappush(heap, int(op1))
        elif heap :
            if op1 == "1" : # 최댓값 삭제
                heap = heapq.nlargest(len(heap), heap)[1:]
                heapq.heapify(heap)
            elif op1 == "-1" :
                heapq.heappop(heap)
    
    if heap :
        min_value = heapq.heappop(heap)
        if heap :
            max_value = heapq.nlargest(1, heap)[0]
        else : 
            max_value = min_value
    else :
        max_value, min_value = 0, 0
    
    return [max_value, min_value]