• 힙에 대한 기본적인 내용은 여기 참고하기!

---

직접 구현하기

파이썬에서 기본으로 제공되는 모듈을 사용하면 편하지만, 동작을 파악하기 위해 직접 구현해보자!

📌 최대 힙을 기준으로 작성했습니다 :)

힙 생성하기

먼저 Heap 클래스를 생성해 초기화한다.
루트의 인덱스 번호를 1로 하기 위해, 리스트의 0번째 자리에 None을 넣어둔다.

class Heap:
    def __init__(self):
        self.heap = []
        self.heap.append(None)

삽입하기

데이터를 삽입할 때, 맨 뒤부터 차례대로 저장한다.
만약 삽입한 데이터가 부모보다 클 경우, 부모와 위치를 바꿔줘야 한다.

# 해당 노드가 부모 노드보다 큰지 비교
def check_swap_up(self, idx):
	# 삽입한 모드의 부모 노드가 없을 경우
    if idx <= 1:
    	return False

	parent_idx = idx // 2

	if self.heap[idx] > self.heap[parent_idx]:
		return True
	else:
    return False

# 데이터 삽입
def insert(self, data):
	self.heap.append(data)
    idx = len(self.heap) - 1

    # check_swap_up() 의 값이 참이라면 부모와 위치 바꾸기
    while self.check_swap_up(idx):
    	parent_idx = idx // 2

        self.heap[idx], self.heap[parent_idx] = self.heap[parent_idx], self.heap[idx]
        idx = parent_idx

	return True

삭제하기

최댓값을 꺼내면 그 자리, 즉 루트 노드가 비어있게 된다.
가장 마지막 노드와 루트 노드의 자리를 바꾼 뒤, 자식 노드와 값을 비교한다.
만약 루트 노드가 자식 노드보다 더 작을 경우, 자식과 위치를 바꾼다.

이때 자식은 없을 수도, 하나만 있을 수도, 둘 다 있을 수도 있다.
하나만 있을 때는 자식과 부모만 비교하면 되지만, 자식이 둘 다 있을 때는 자식끼리 먼저 비교하여 더 큰 자식과 부모를 비교한다.

# 해당 노드가 부모 노드보다 큰지 비교
def check_swap_down(self, idx):
	left_idx = idx * 2
    right_idx = idx * 2 + 1
        
    # 자식 노드가 하나도 없을 경우
    if left_idx >= len(self.heap):
    	return False
        
	# 왼쪽 자식 노드만 있을 경우
    elif right_idx >= len(self.heap):
    	if self.heap[left_idx] > self.heap[idx]:
        	self.flag = 1
            return True
		else:
        	return False
        
	# 자식 노드가 모두 있을 경우
    else:
    	if self.heap[left_idx] > self.heap[right_idx]:
        	if self.heap[left_idx] > self.heap[idx]:
            	self.flag = 1
                return True
			else:
            	return False
		else:
        	if self.heap[right_idx] > self.heap[idx]:
            	self.flag = 2
                return True
			else:
            	return False

# 데이터 삭제
def pop(self):
	if len(self.heap) <= 1:
    	return None
        
	max = self.heap[1]
    self.heap[1] = self.heap[-1]
    del self.heap[-1]
    idx = 1

    # 0 = False, 1 = (왼쪽 자식과 swap), 2 = (오른쪽 자식과 swap)
    self.flag = 0 

    while self.check_swap_down(idx):
    	left_idx = idx * 2
        right_idx = idx * 2 + 1

        if self.flag == 1:
        	self.heap[idx], self.heap[left_idx] = self.heap[left_idx], self.heap[idx]
            idx = left_idx
		elif self.flag == 2:
        	self.heap[idx], self.heap[right_idx] = self.heap[right_idx], self.heap[idx]
            idx = right_idx
	return max

파이썬 heapq 모듈 사용하기

파이썬에서 기본적으로 제공하는 heapq 모듈은 우선순위 큐(Priority Queue) 알고리즘을 제공한다.

파이썬의 리스트를 사용해 인덱스 0부터 시작해 k번째 원소가 항상 자식 원소(2k+1, 2k+2)보다 작거나 같은 최소 힙의 형태로 정렬한다.

import

heapq 모듈은 파이썬의 내장 모듈이기 때문에 따로 설치하지 않아도 된다.

import heapq

힙 생성하기

heapq 모듈은 파이썬의 리스트를 최소 힙 형태로 정렬하기 때문에 빈 리스트를 생성한 뒤, 모듈 함수를 호출할 때마다 생성한 리스트를 인자 값으로 넘겨야 한다.

즉, 추가/삭제 작업만 해도 알아서 정렬해준단 얘기!

heap = []

삽입하기

모듈의 heappush() 함수를 사용해 원소를 삽입할 수 있다.
첫 번째 인자로는 대상 리스트를, 두 번째 인자로는 삽입할 값을 전달한다.

heapq.heappush(heap, 10)
heapq.heappush(heap, 6)
heapq.heappush(heap, 13)
heapq.heappush(heap, 5)

print(heap)

[5, 6, 13, 10]

삭제하기

모듈의 heappop() 함수를 사용해 원소를 삭제할 수 있다.
대상 리스트를 인자로 넘기면, 최솟값을 삭제한 뒤에 반환한다.

print(heapq.heappop(heap))
print(heap)

5
[6, 13, 10]

이때, 삭제하지 않고 최솟값을 출력하기 위해서는 다음과 같이 작성한다.

print(heap[0])

6

그렇다면 heap[1]에는 두 번째로 작은 원소가 들어있을까??

답은 ❌

앞에서 출력한 결과만 봐도 알 수 있다.

[6, 13, 10]

작은 순서대로 정렬이 되어 있다면 10이 13보다 앞에 위치해야 하는데, 그렇지 않다.

최소 힙은 최솟값을 빠르게 찾기 위한 알고리즘이지, 작은 순서대로 정렬하는 알고리즘이 아니다.
때문에 작은 순서대로 정렬되어 있을 것이라고 생각하지 말자!!

리스트를 힙으로 변환하기

이미 원소가 삽입되어 있는 리스트를 heapq 모듈을 사용해 힙으로 변환하는 것이 가능하다.
이때는 heapify() 함수를 사용하고, 인자로 대상 리스트를 전달한다.

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

print(heap)

[1, 2, 4, 3, 7]

응용하기

백준 알고리즘 1927번

여기에서 문제 확인하기!

널리 잘 알려진 자료구조 중 최소 힙이 있다. 최소 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.

배열에 자연수 x를 넣는다.
배열에서 가장 작은 값을 출력하고, 그 값을 배열에서 제거한다.
프로그램은 처음에 비어있는 배열에서 시작하게 된다.

입력

첫째 줄에 연산의 개수 N(1 ≤ N ≤ 100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 작은 값을 출력하고 그 값을 배열에서 제거하는 경우이다. x는 231보다 작은 자연수 또는 0이고, 음의 정수는 입력으로 주어지지 않는다.

출력

입력에서 0이 주어진 횟수만큼 답을 출력한다. 만약 배열이 비어 있는 경우인데 가장 작은 값을 출력하라고 한 경우에는 0을 출력하면 된다.

예제 입력1

9
0
12345678
1
2
0
0
0
0
32

0
1
2
12345678
0

---

작성한 코드

앞에서 설명한 heapq 모듈만 기억한다면 쉽게 풀 수 있는 문제!

import heapq
import sys

N = int(sys.stdin.readline())
heap = list()

for i in range(N):
    x = int(sys.stdin.readline())

    if x==0:
        if not heap:
            print(0)
        else:
            print(heapq.heappop(heap))
    
    else:
        heapq.heappush(heap, x)