[파이썬 ] 우선순위 큐(우선 순위 대기열)를 위한 heapq모듈 활용법

우선순위 큐는 우선 순위가 가장 높은 자료(데이터)를 가장 먼저 꺼낼 수 있는 자료구조다. 배열을 사용하면 직접 구현하기는 어렵지 않지만, 파이썬에서는 heapq 라는 내장(내장)모듈로 제공되기 때문에, 내장 모듈의 미국 석유 학회기능만으로 충분히 활용 가능할때는 내장 모듈을 사용하는 것이 좋다.
이 글에서 힙은 이진힙을 의미하고, 우선순위 큐는 개별의 값을 가지는 힙이므로 두 단어를 번갈아 사용하겠다

기본 사용법

1우선 순위 큐의 생성 및 원소 삽입heapq.heappush 를 사용해 우선 순위 큐의 원소를 삽입할 수 있다. 첫번째 인자는 힙으로 사용할 리스트이고, 두번째 인자는 삽입할 데이터이다.heapq.heappush(heap, item)삽입별 시간 복잡도는 O(대수 n)이다. n개의 원소를 삽입한다면 로그의 시간 복잡도를 가진다.
다음 예제를 보자.

import heapq

h = []
heapq.heappush(h, (3, "Go to home"))
heapq.heappush(h, (10, "Do not study"))
heapq.heappush(h, (1, "Enjoy!"))
heapq.heappush(h, (4, "Eat!"))
heapq.heappush(h, (7, "Pray!"))

print(h)

결과는 배열로 구현된 힙을 보여준다.

[(1, 'Enjoy!'), (4, 'Eat!'), (3, 'Go to home'), (10, 'Do not study'), (7, 'Pray!')]

2원소 꺼내기
원소 꺼내기는 heapq.heappop 으로 구현하면 된다. 각 꺼내기 별 시간 복잡도는 O(대수 n)이다.

import heapq

h = []
heapq.heappush(h, (3, "Go to home"))
heapq.heappush(h, (10, "Do not study"))
heapq.heappush(h, (1, "Enjoy!"))
heapq.heappush(h, (4, "Eat!"))
heapq.heappush(h, (7, "Pray!"))

count = 1
while h:
    print("%dth : %s" % (count, heapq.heappop(h)))
    count += 1

우선순위 순서대로 나온 결과를 확인할 수 있다.

1th : (1, 'Enjoy!')
2th : (3, 'Go to home')
3th : (4, 'Eat!')
4th : (7, 'Pray!')
5th : (10, 'Do not study')

셋.배열로부터 우선순위 큐 만들기
1번의 방법으로 n개의 요소를 가진 우선순위 큐를 만들었을 때, 시간 복잡도는 로그이다. 그러나 배열로부터 힙을 만드는 최적 알고리즘의 시간 복잡도는 O(n)으로 알려져 있다 . 파이썬에서는 O(n)의 시간으로 배열을 힙으로 만들 수 있는 heapq.heapfy 을 제공한다. 주의할 점은 인자로 넣은 배열 자체가 힙으로 바뀐다는 점이다.

import heapq

h = [(3, "Go to home"), (10, "Do not study"), (1, "Enjoy!"), (4, "Eat!"), (7, "Pray!")]
heapq.heapify(h)

count = 1
while h:
    print("%dth : %s" % (count, heapq.heappop(h)))
    count += 1

결과는 1번으로부터 생성한 우선순위 큐와 동일하다.

1th : (1, 'Enjoy!')
2th : (3, 'Go to home')
3th : (4, 'Eat!')
4th : (7, 'Pray!')
5th : (10, 'Do not study')

heapq의 한계점
눈치 챘을 수도 있겠지만, heapq 모듈은 최소 힙(최소 더미)밖에 지원을 안한다. 최대 힙(최대 더미)을 사용할 수 있는 방법은 없을까? heapq 에서 공식적으로 지원해주는 기능은 없다. 검색을 해보면, heapq._heapfy_max 나 heapq._heappop_max 를 사용하는 방법도 있지만, 밀다를 지원해주지 않기 때문에 반쪽짜리 기능이고, 보호 방법라서, 활용단어참조의 버전이 바뀌면서 이음매가 변경될 수 있는 위험도 있다.
유일한 방법은 키 값을 변환하여 삽입하는 수 밖에 없다. 위에서 사용한 예제를 기준으로 최대 힙을 구현한다고 하면 다음과 같다

import heapq

h = [(3, "Go to home"), (10, "Do not study"), (1, "Enjoy!"), (4, "Eat!"), (7, "Pray!")]
max_h = [(-ele[0], ele[1]) for ele in h]
heapq.heapify(max_h)


count = 1
while h:
    print("%dth : %s" % (count, heapq.heappop(max_h)))
    count += 1

결과는 다음과 같다.

1th : (-10, 'Do not study')
2th : (-7, 'Pray!')
3th : (-4, 'Eat!')
4th : (-3, 'Go to home')
5th : (-1, 'Enjoy!')

모든 요소의 값을 변경하는데 드는 시간복잡도는 O(n)이고, 모든 힙 요소를 얻는데 드는 시간 복잡도는 로그이므로, 요소를 모두 돌면서 키를 바꾼다 해도, 전체 시간복잡도에는 영향을 주지 않는다.
그러나 이 방법은 원본 데이터를 건드릴 수 밖에 없는 점과, heapq.push 를 사용해 추가로 데이터를 삽입할 때마다 데이터를 변경해줘야 하고, 추출했을 때도 원본 데이터와 다르다는 단점이 있다.

해결 방법 - 최대 힙 모듈 구현
직접 데이터를 변경하여 넣어줄 때 생기는 단점을 극복할 수 있는 방법중에, 기존 heapq 모듈을 래핑하여, 최대 힙을 구현하는 방법도 있다. 파이썬의 __lt__ 내장 메소드를 사용하면, 어렵지 않게 구현할 수 있다.

# maxheapq.py
import heapq

class _ReverseLessThan(object):
    def __init__(self, value):
        self.value = value

    def __lt__(self, other):
        return self.value > other.value

    def __repr__(self):
        return str(self.value)


def heappush(heap, item):
    reverse_item = _ReverseLessThan(item)
    heapq.heappush(heap, reverse_item)


def heappop(heap):
    reverse_item = heapq.heappop(heap)
    return reverse_item.value


def heapify(lst):
    for i, ele in enumerate(lst):
        lst[i] = _ReverseLessThan(ele)
    heapq.heapify(lst)

_ReverseLessThan 클래스는 생성자( __init__ )에서 데이터를 인자로 받아서, value재산로 저장한 후, __lt__ 메서드의 결과를 value재산기준으로 역으로 반환하여, 이 클래스의 인스턴스로 조금비교 연산 시, 비교 연산 결과가 반대로 나오도록 해준다.heappush , heappop , heapify 함수는 _ReverseLessThan 을 이용해, heapq 모듈의 함수를 래핑하여 구현하였다. heapq 모듈은 최소 힙을 지원하므로, 조금연산의 결과가 반대로 나오는 인스턴스를 데이터로 넣으면, 최대힙이 구현된다.
주의할 점은 삽입 연산인 heappush , heapify 에서는 _ReverseLessThan 클래스로 데이터를 래핑해주고 있으므로, 추출 연산인 heappop 에서는 언래핑 해주어야 한다.
각 함수의 전체 시간 복잡도는 heapq 모듈의 함수와 동일하다.
이 모듈을 사용해, 최대 힙을 만들어보자.

from maxheapq import heappush, heappop, heapify

# heappush로 생성하기
h1 = []

heappush(h1, (3, "Go to home"))
heappush(h1, (10, "Do not study"))
heappush(h1, (1, "Enjoy!"))
heappush(h1, (4, "Eat!"))
heappush(h1, (7, "Pray!"))
heappush(h1, (4, "Dup Eat!"))

print('>> h1 생성 결과')
print(h1)

count = 1
while h1:
    print("%dth : %s" % (count, heappop(h1)))
    count += 1


# heapify로 생성하기
h2 = [(3, "Go to home"), (10, "Do not study"), (1, "Enjoy!"), (4, "Eat!"), (7, "Pray!"), (4, "Dup Eat!")]

heapify(h2)

print('>> h2 생성 결과')
print(h2)

count = 1
while h2:
    print("%dth : %s" % (count, heappop(h2)))
    count += 1

결과를 보면 최대 힙이 정상적으로 생성된 것을 볼 수 있다.

>> h1 생성 결과
[(10, 'Do not study'), (7, 'Pray!'), (4, 'Dup Eat!'), (3, 'Go to home'), (4, 'Eat!'), (1, 'Enjoy!')]
1th : (10, 'Do not study')
2th : (7, 'Pray!')
3th : (4, 'Eat!')
4th : (4, 'Dup Eat!')
5th : (3, 'Go to home')
6th : (1, 'Enjoy!')
>> h2 생성 결과
[(10, 'Do not study'), (7, 'Pray!'), (4, 'Dup Eat!'), (4, 'Eat!'), (3, 'Go to home'), (1, 'Enjoy!')]
1th : (10, 'Do not study')
2th : (7, 'Pray!')
3th : (4, 'Eat!')
4th : (4, 'Dup Eat!')
5th : (3, 'Go to home')
6th : (1, 'Enjoy!')

지금까지 파이썬의 최소 힙 모듈인 heapq 를 활용하는 간단한 방법과, heapq 에서 제공해주지 않는 최대 힙을 직접 구현하여 사용할 수 있는 방법을 알아보았다. 실제 업무에서 이진힙을 사용하는 경우는 많지는 않지만, 혹시나 활용할 기회가 생긴다면 이글에 있는 내용이 도움이 되기를 기대해본다.