Python의 heapq는 sort보다 빠른가..?

코딩테스트 문제를 풀던 중, 다른 사람의 풀이에서 정렬을 용이하게 하기 위해 heapq를 사용하는것을 보았다.

필자는 heapq를 사용하면 빠른가 싶어서 heapq를 사용했으나, 여전히 시간초과가 되는 것을 보았다.

이에 heapq는 정말로 빠른가 하는 것에 의문이 들었다.

 

heapq가 무엇이고 어떻게 사용하는지는 아래 링크의 글을 통해 알게되었다.

(글의 나머지 부분에서 설명할 heappush, heappop에 대한 설명 나옴)

파이썬의 heapq 모듈로 힙 자료구조 사용하기 | Engineering Blog by Dale Seo

 

heappush와 sort의 실행 시간을 비교하기 위해 아래와 같은 코드를 만들어 테스트해보았다.

 

import random
from heapq import heappush, heappop

n = 10_000_000
list = []

do_sort = False
# do_sort = True

if do_sort:
	for i in range(n):
		list.append(random.randint(0, 100))
	list.sort()

else:
	for i in range(n):
		heappush(list, random.randint(0, 100))

필자의 테스트 환경에서는 n = 10_000_000 일 때 sort 방법과 heappush 방법이 모두 15초 정도가 걸렸다.

n을 20_000_000로 높여서 테스트해보았는데, 둘 다 28~32초 사이에서 왔다갔다한다.

저 범위를 넘어가면 모르겠지만, 저 범위에서까지는 둘의 성능 차이를 거의 느낄 수 없었다.

 

이번엔 pop 까지 테스트해보았다.

 

import random
from heapq import heappush, heappop

n = 10_000_000
list = []

do_sort = False
# do_sort = True

if do_sort:
	for i in range(n):
		list.append(random.randint(0, 100))
	list.sort()

	for i in range(n):
		list.pop()
else:
	for i in range(n):
		heappush(list, random.randint(0, 100))

	for i in range(n):
		heappop(list)

 

sort & pop이 18초, heapq가 25초로, 확실히 sort & pop이 훨씬 빨랐다.

 

참고로 heapq를 할 수 있는 리스트가 따로 있는게 아니라, 그냥 list를 heapq 함수에 전달하면 되는거라

필자는 heappush를 이용하지 않아도 뒤에서부터 일반 pop하면 된다고 착각했지만,

위 링크의 글에도 있듯이 이진트리 구조를 사용하기 때문에, 그렇게 하면 순서가 꼬여서 제대로 쓸 수 없다.

 

혹시 필자의 테스트 환경에 문제가 있는건지, 진짜로 heapq가 느린건지 싶어서 검색을 해보았다.

아래 글에 heappush와 heappop의 시간 복잡도에 대한 정보가 있어 시간 계산에 참고했다.

[Python] heapq 모듈 (velog.io)

 

위 글에 따르면 heappush는 O(logN) 시간이 걸린다고 한다.

(아마 정렬된 리스트에서 새 원소가 들어갈 위치를 binary search로 계산하기 때문이 아닐까 싶다.)

원소가 N개 이미 들어있을 때 새 원소를 추가하는 시간이 대략 logN이라는 뜻일테니,

비어있는 원소에 0부터 N-1까지 원소를 추가하는데 걸리는 시간은 Sigma(logN) -> NlogN이 된다.

(Sigma(logN) 설명 참고: https://math.stackexchange.com/a/228748)

 

원소 추가 후 sort의 경우는,

N(원소 추가) + NlogN(정렬) 만큼 걸린다. N이 많이 큰 경우는 O(NlogN)이 된다.

 

pop의 경우는 heappop이 훨씬 느리다.

역시 위 링크의 글에 따르면 heappop은 O(logN)이 걸린다고 하니, (이유는 모른다..)

O(1)이 걸리는 일반 pop보다 훨씬 느리다.

게다가 N번 수행하는거면 O(NlogN) vs O(N)으로 일반 pop이 훨씬 빠르다.

 

그렇다면 heapq는 언제 사용하면 좋을까?

코딩테스트 상황처럼 빈 리스트에 새 원소들을 추가하며 최종적으로 정렬된 리스트를 얻어야 하는 상황보다는

새로운 데이터가 추가되어도 항상 정렬 상태를 유지해야 하는 상황이거나,

이미 정렬되어있는 리스트에 새 원소를 추가하는 경우에 유용하게 쓰일 수 있을 것 같다.

 

 

결론: heapq는 새로운 데이터가 추가되어도 항상 정렬된 상태를 유지하게 하거나, 이미 정렬되어 있는 상태를 이용하기 위해 사용하는것이지,

정렬을 빠르게 하는 것과는 관련이 없다.

데이터를 모두 입력받은 뒤 사용은 나중에 할 것이라면 heapq를 쓰느니 sort()를 쓰는게 더 낫다.