[Python3] heapq 사용법

민갱·2023년 7월 10일

Python

목록 보기

10/11

자료구조 힙을 사용하기 위해서 heapq의 사용법과 관련 함수들에 대해 공부해보자.

1. 모듈 임포트

heapq 모듈은 python 내장 모듈이기 때문에 간단히 import 하여 사용한다.

from heapq import *
or
import heapq

2. 최소 힙 생성

heapq 모듈은 파이썬의 리스트를 마치 최소 힙처럼 다룰 수 있도록 도와준다.
따라서 heapq 모듈을 통해서 리스트에 원소를 추가, 삭제하면 그 리스트가 최소힙이 된다.

heap = [] // heapq를 사용해 최소 힙으로 다룰 리스트

3. 힙에 원소 추가

heapq 모듈의 heappush() 함수를 이용하여 힙에 원소를 추가한다.
첫번째 인자는 원소를 추가할 대상 리스트이고, 두번째 인자는 추가할 요소 이다.

heapq.heappush(heap,4)
heapq.heappush(heap,1)
heapq.heappush(heap,7)
heapq.heappush(heap,3)
print(heap)

==>
[1,3,7,4]

가장 작은 1이 인덱스 0에 위치하며, 인덱스 1(=k)에 위치한 3은 인덱스 3(=2k+1)에 위치한 4보다 작으므로 힙의 공식에 만족한다.
내부적으로 이진 트리에 원소를 추가하는 heappush() 함수는 O(logN)의 시간 복잡도를 갖는다.
heapq의 heappush로 원소를 추가하면 자체적으로 heap 공식에 맞게 요소를 넣어준다.

4. 힙에서 원소 삭제

heapq 모듈의 heappop() 함수를 이용하여 힙에서 원소를 삭제할 수 있다. 원소를 삭제할 대상 리스트를 인자로 넘기면, 가장 작은 원소(heap[0])를 삭제 후 그 값을 리턴한다.


print(heapq.heappop(heap))
print(heap)
==>
1
[3,4,7]

가장 작았던 1이 삭제되어 리턴되었고, 그 다음으로 작었던 3이 인덱스 0으로 올라왔다.

print(heapq.heappop(heap))
print(heap)
==>
3
[4,7]

가장 작았던 3이 삭제되어 리턴되고, 그다음으로 작았던 4가 인덱스 0으로 올라왔다. 내부적으로 이진 트리로 부터 원소를 삭제하는 heappop() 함수 역시 O(logN)의 시간복잡도를 갖는다.

5. 최소값 삭제하지 않고 얻기 (peek)

힙에서 최소값을 삭제하지 않고 단순히 읽기만 하려면 일반적으로 리스트의 첫번째 원소에 접근하듯이 인덱스를 통해 접근하면 된다.

print(heap[0])
==>
4

여기서 주의사항은 인덱스 0에 가장 작은 원소가 있다고 해서, 인덱스 1에 두번째 작은 원소, 인덱스 2에 세번째 작은 원소가 있다는 보장은 없다는 것이다.
왜냐하면 힙은 heappop() 함수를 호출하여 원소를 삭제할 때마다 이진 트리의 재배치를 통해 매번 새로운 최소값을 인덱스 0에 위치시키기 때문이다.

따라서 두번째로 작은 원소를 얻으려면 바로 heap[1]으로 접근하면 안되고, 반드시 heappop()을 통해 가장 작은 원소를 삭제 후에 heap[0]를 통해 새로운 최소값에 접근해야 한다.

6. 기존 리스트를 힙으로 변환

이미 원소가 들어있는 리스트를 힙으로 만드려면 heapq 모듈의 heapify() 함수를 이용하면 된다.

heap = [4,1,7,3,8,5]
heapq.heapify(heap)
print(heap)
==>
[1,3,5,4,8,7]

heapify() 함수에 리스트를 인자로 넘기면 리스트 내부의 우너소들의 위에서 다룬 힙 구조에 맞게 재배치되며 최소값이 0번째 인덱스에 위치한다.
즉, 비어있는 리스트를 생성한 후 heappush() 함수로 원소를 하나씩 추가한 효과가 난다.
heapify() 함수의 성능은 인자로 넘기는 리스트 원소수에 비례한다. O(logN) 의 시간복잡도를 갖는다.

응용 1. 최대 힙

heapq 모듈은 최소 힙(min heap)만을 만들기 때문에, 최대 힙(max heap)으로 활용하려면 약간의 요령이 필요하다. 바로 힙에 튜플(tuple) 원소를 추가하거나 삭제하면, 튜플 내에서 맨 앞에 있는 값을 기준으로 최소 힙이 구성되는 원리를 이용하는 것이다.
따라서 최대 힙을 만드려면 각 값에 대한 우선순위를 구한 후, (우선 순위, 값) 구조의 튜플을 힙에 추가하거나 삭제하면 된다.
힙에서 값을 읽어올 때는 각 튜플에서 인덱스 1에 있는 값을 취하면 된다.(우선 순위는 관심이 없다.)

import heapq

nums = [4, 1, 7, 3, 8, 5]
heap = []

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

while heap:
  print(heapq.heappop(heap)[1])  # index 1
  
==>
 
8
7
5
4
3
1

응용 2. K번째 최소값/최대값

최소 힙이나 최대 힙을 사용하면 K번째 최소값이나 최대값을 효율적으로 구할 수 있다.


import heapq

# min heap
def kth_smallest(nums,k):
	heap = []
    for num in nums:
    	heapq.heappush(heap, num)
        
    kth_min = None
    for _ in range(k):
    	kth_min = heapq.heappop(heap)
    return kth_min
print(kth_smalles([4,1,7,3,8,5],3))
==>
4
# max heap

def kth_biggest(nums,k):
	heap = []
    for num in nums:
    	heapq.heappush(heap, -(num,num))
        
    kth_max = None
    for _ in range(k):
    	kth_max = heapq.heappop(heap)[1]
    return kth_max
print(kth_biggest([4,1,7,3,8,5],3))
==>
5

K번째 최소값을 구하기 위해서는 주어진 배열로 힙을 만든 후, heappop() 함수를 K번 호출하면 된다.

응용3. 힙 정렬

힙 정렬(heap sort)은 힙 자료구조의 성질을 이용한 대표적인 정렬 알고리즘이다.


import heapq

def heap_sort(nums):
	heap = []
    for num in nums:
    	heapq.heappush(heap,num)
   
    sorted_nums =[]
    while heap:
    	sorted_nums.append(heapq.heappop(heap))
   	return sorted_nums

print(heap_sort([4,1,7,3,8,5]))
==>
[1,3,4,5,7,8]

heapq 함수


# heapq 함수 정리
import heapq

heap = []
# 삽입
heapq.heappush(heap, value)

# 삭제
heapq.heappop(heap)

# peek
heap[0]

# 리스트를 heap으로 변환
heapq.heapify(heap)