[파이썬 ] 우선순위 큐(priority queue)를 위한 heapq 모듈 활용법
Tags: python, heapq, priorityqueue, binaryheap
Published: 2021-04-14 12:16:29

우선순위 큐는 우선 순위가 가장 높은 자료(data)를 가장 먼저 꺼낼 수 있는 자료구조다. 배열을 사용하면 직접 구현하기는 어렵지 않지만, 파이썬에서는 `heapq`라는 내장(built-in) 모듈로 제공되기 때문에, 내장 모듈의 API 기능만으로 충분히 활용 가능할때는 내장 모듈을 사용하는 것이 좋다.

*이 글에서 힙은 이진힙을 의미하고, 우선순위 큐는 개별의 값을 가지는 힙이므로 두 단어를 번갈아 사용하겠다*
</br>
</br>


# 기본 사용법
</br>
</br>
## 1. 우선 순위 큐의 생성 및 원소 삽입
`heapq.heappush`를 사용해 우선 순위 큐의 원소를 삽입할 수 있다. 첫번째 인자는 힙으로 사용할 list 이고, 두번째 인자는 삽입할 데이터이다.

`heapq.heappush(heap, item)`

삽입별 시간 복잡도는 O(log n)이다. n개의 원소를 삽입한다면 O(n log n)의 시간 복잡도를 가진다.

다음 예제를 보자.

```python
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!')]
```
</br>
</br>

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

```python
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')
```
</br>
</br>

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

```python
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')
```
</br>
</br>

# heapq의 한계점
 눈치 챘을 수도 있겠지만, `heapq` 모듈은 최소 힙(min heap) 밖에 지원을 안한다. 최대 힙(max heap)을 사용할 수 있는 방법은 없을까? `heapq`에서 공식적으로 지원해주는 기능은 없다. 검색을 해보면, `heapq._heapfy_max` 나 `heapq._heappop_max`를 사용하는 방법도 있지만, push를 지원해주지 않기 때문에 반쪽짜리 기능이고, protected method라서, python의 버전이 바뀌면서 interface가 변경될 수 있는 위험도 있다.

유일한 방법은 키 값을 변환하여 삽입하는 수 밖에 없다. 위에서 사용한 예제를 기준으로 최대 힙을 구현한다고 하면 다음과 같다

```python
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)이고, 모든 힙 요소를 얻는데 드는 시간 복잡도는 O(n log n)이므로, 요소를 모두 돌면서 키를 바꾼다 해도, 전체 시간복잡도에는 영향을 주지 않는다.

그러나 이 방법은 원본 데이터를 건드릴 수 밖에 없는 점과, `heapq.push`를 사용해 추가로 데이터를 삽입할 때마다 데이터를 변경해줘야 하고, 추출했을 때도 원본 데이터와 다르다는 단점이 있다.
</br>
</br>

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


```python
# 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` property로 저장한 후, `__lt__` 메서드의 결과를 `value` property 기준으로 역으로 반환하여, 이 클래스의 인스턴스로 less 비교 연산 시, 비교 연산 결과가 반대로 나오도록 해준다.

`heappush`, `heappop`, `heapify` 함수는 `_ReverseLessThan`을 이용해, `heapq` 모듈의 함수를 래핑하여 구현하였다. `heapq` 모듈은 최소 힙을 지원하므로, less 연산의 결과가 반대로 나오는 인스턴스를 데이터로 넣으면, 최대힙이 구현된다.

주의할 점은 삽입 연산인 `heappush`, `heapify` 에서는 `_ReverseLessThan` 클래스로 데이터를 래핑해주고 있으므로, 추출 연산인 `heappop`에서는 언래핑 해주어야 한다.

각 함수의 전체 시간 복잡도는 `heapq` 모듈의 함수와 동일하다.


이 모듈을 사용해, 최대 힙을 만들어보자.

```python
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`에서 제공해주지 않는 최대 힙을 직접 구현하여 사용할 수 있는 방법을 알아보았다. 실제 업무에서 이진힙을 사용하는 경우는 많지는 않지만, 혹시나 활용할 기회가 생긴다면 이글에 있는 내용이 도움이 되기를 기대해본다.