힙(Heap) 자료구조

     1  ---> root
   /   \
  3     5
 / \   /
4   8 7

힙은 완전 이진 트리 자료구조의 일종으로, 힙에서는 항상 루트 노드를 제거한다.

파이썬에서는 heapq모듈을 임포트하여 사용가능하며, heapq는 최소 힙(min heap) 자료구조를 제공한다.

모듈 임포트

import heapq

힙 사용

1. 빈 리스트를 생성하여 힙처럼 다루는 방법

heap = []
heapq.heappush(heap, 4) # 힙에 원소 추가하는 함수
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
heapq.heappush(heap, 3)
print(heap)

[1, 3, 7, 4]	# 실행 결과

heappush() 함수를 이용하여 힙에 원소를 추가할 수 있다.

2. 기존에 있던 리스트를 Heapify를 이용해 heap으로 만드는 방법

heap = [3, 4, 7, 2, 1]
heapq.heapify(heap)
print(heap)

[1, 2, 7, 3, 4] # 실행 결과

heapify() 함수를 이용하면 리스트의 원소들을 힙 구조에 맞게 재배치하게된다.

3. 힙에서 원소 삭제

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

1				# 실행 결과
[2, 3, 7, 4]

heappop() 함수를 이용하여 원소를 삭제할 수 있고, 가장 작은 원소를 삭제한 후에 리턴한다.

최대 힙(max heap)

heapq 모듈은 최소 힙(mini heap)을 기준으로 동작하기 때문에 최대 힙(max heap)으로 활용하려면 모든 값에 -를 붙여 최대 값이 최소 값이 되도록 하면 된다.

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

예제

가장 작은 두 수를 계속 더해 주는데, 이때 더한 값을 다시 포함시켜서 계속 더해준다.

from sys import stdin
import heapq

n = int(stdin.readline().rstrip('\n'))
data = []
for _ in range(n):
    data.append(int(stdin.readline().rstrip('\n')))
heapq.heapify(data)
ans = 0

while len(data) != 1:
    num1 = heapq.heappop(data)
    num2 = heapq.heappop(data)
    sum = num1 + num2
    ans += sum
    heapq.heappush(data, sum)

print(ans)

입력 값을 모두 data라는 리스트에 넣어주고 heap구조로 바꿔준다.
그리고 다 더해서 최종 값이 나올때 까지 가장 작은 수를 빼서 더하고 합한 값만 다시 data에 집어넣는 것을 반복한다.

출처: heapq 모듈 사용법, 자료구조: 우선순위 큐