Golang으로 Min Heap을 구현해보자

Photo by Markus Spiske on Unsplash

 

LeetCode 문제를 풀다가 min heap이 필요한 문제를 만난김에, Go는 쉽게 쓸 수 있는 패키지는 제공하지 않기에 (오히려 그 덕분에) min heap을 들여다보고 공부해보게 되었다.

(참고: LeetCode problem: https://leetcode.com/contest/weekly-contest-237/problems/single-threaded-cpu/ )

 

 

Min Heap이란 무엇인가?

이진 트리인데 parent는 아래의 child 노드들보다 작은 값을 가진다. 실제 구현은 슬라이스의 형태이다.

 

< 이미지 출처: https://www.wikiwand.com/en/Min-max_heap >

 

아래 이미지는 Max Heap이지만,  대괄호 속 노란색 숫자가 슬라이스의 인덱스라고 보면 이해가 쉬울 것이다.

 

< 이미지 출처: https://youtu.be/3DYIgTC4T1o 캡처 >

 

구현해야 할 것들

Push 와 up

1) heap 으로 하나가 들어오면 그걸 가장 끝의 leaf에 추가한다. 위 그림으로 보면 12번 인덱스 위치가 될 것이다.

2) 그리고 min heap이라면 부모와 비교하며 자신이 작다면, 계속해서 거슬러 올라가게 된다.

Pop 과 down

1) 우선 0번 인덱스의 값을 빼놓고 나중에 return 해준다. min heap 이라면 0번 인덱스의 값이 항상 가장 작은 값이다.

2) 그리고는 맨 뒤 인덱스의 값과 0번 인덱스 값을 swap 한 다음에 전체 길이를 하나 줄인다.

3) 마지막으로 0번 인덱스의 값을 자식들과 비교해가며 자신이 더 크면 계속 내려가게 한다.

swap 과 less

less를 통해 노드간의 비교 규칙을 정해두고

swap으로 두 인덱스의 값을 바꿔치게 하는 것이다.

Init

이건 꼭 필요한 것은 아니지만, 이미 주어진 슬라이스가 있을때에, 이것을 heap을 만족하게 만들어준다.

방법은 간단하다.

 

1) 부모중에서 가장 큰 인덱스를 찾는다. 가장 끝 leaf의 부모이다.

2) 가장 큰 인덱스의 부모값부터 인덱스를 하나씩 줄여가며 down을 해주면 된다.

꼭 필요하지는 않은, 부모, 자식 인덱스 계산기

특정 인덱스의 부모 인덱스, 왼쪽, 오른쪽 자식 인덱스를 계산하는 간단한 함수가 있으면 되겠다.

 

실제 코드 구현 분석

여기서 노드는 두 개의 int 값을 가지는 슬라이스이다. 슬라이스의 첫번째 값으로 먼저 비교하고, 같으면 두 번째 값이 작은 순으로 heap 정렬을 한다. 그냥 int 하나만 가지는 예제였으면 좀 더 이해하기 좋았었겠지만 오히려 어떠한 값이라도 heap의 값이 될 수 있다는 좋은 예가 될 수도 있겠다.

 

구현 예제: https://play.golang.org/p/cjYXpQjPiZY

Push 와 up

힙 구현은 결국 슬라이스이다. 어떤 값을 가진 노드들을 슬라이스안에 넣는 것이다.

여기서는 슬라이스 값을 담는 슬라이스이다.

 

Push 메서드는 슬라이스의 맨 뒤(=가장 마지막 leaf이다)에 추가되는 값을 넣은 다음에, up 메서드를 통해 재귀적으로 부모 인덱스의 값과 비교하여 바꿔나간다. (참고: 유튜브 예제에서는 구조체 안에 슬라이스를 넣었는데 구현상의 깔끔함에 있어서 그게 더 나은 것 같다. *를 일일이 붙이지 않아도 된다)

 

type MinHeap [][]int

func (h *MinHeap) Push(vals []int) {
	*h = append(*h, vals)
	h.up(len(*h) - 1)
	fmt.Printf("%v Pushed, heap is now: %v\n", vals, h)
}

func (h *MinHeap) up(idx int) {
	for h.less(idx, parentOf(idx)) {
		h.swap(parentOf(idx), idx)
		idx = parentOf(idx)
	}
}

 

Pop 과 down

Pop 메서드는 우선 인덱스 0의 값을 미리 빼두고, 맨 뒤의 값을 인덱스 0의 값과 바꾼뒤 슬라이스 범위를 하나 줄여서 결국은 원래 인덱스 0 값을 heap에서 없애준다. 그리고 인덱스 0의 값을 down 메서드를 이용해 자식들과 비교해가며 원래 자리를 찾아가는 동시에 heap이 제대로 구성되게 해준다

 

func (h *MinHeap) Pop() []int {
	if len(*h) <= 0 {
		fmt.Println("nothing to pop from the heap")
		return nil
	}
	popped := (*h)[0]
	(*h)[0] = (*h)[len(*h)-1]
	(*h) = (*h)[:len(*h)-1]
	h.down(0)
	return popped
}

func (h *MinHeap) down(idx int) {
	lastIdx := len(*h) - 1
	l, r := leftOf(idx), rightOf(idx)

	var childToCompare int
	for l <= lastIdx {
		if l == lastIdx {
			childToCompare = l
		} else if h.less(l, r) {
			childToCompare = l
		} else {
			childToCompare = r
		}

	if h.less(idx, childToCompare) {
		return
	}
	h.swap(idx, childToCompare)
	idx = childToCompare
	l, r = leftOf(idx), rightOf(idx)
	}
}

swap 과 less

swap 메서드는 단순히 두 원소의 위치를 바꿔주는 것이다.

less 메서드가 중요하다면 중요한데, 두 인덱스의 값을 바꿀지 여부를 결정하는 것이다. 여기의 부등호들의 방향을 바꾸면 min heap은 max heap 이 될 것이다.

func (h *MinHeap) swap(i, j int) {
	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
}

func (h *MinHeap) less(i, j int) bool {
	// we assume node is [2]int
	if (*h)[i][0] != (*h)[j][0] {
		return (*h)[i][0] < (*h)[j][0]
	}
	return (*h)[i][1] < (*h)[j][1]
}

Init

init 메서드는 이미 들어와있는 슬라이스를 heap 데이터 구조로 만들어준다.

원리는 간단하다. 가장 마지막 원소의 부모는 가장 마지막 부모이다. 이 부모의 값 부터 0번 인덱스까지 down 메서드를 이용해서 제자리를 찾아가게 해준다.

func (h *MinHeap) Init() {
	for i := parentOf(len(*h) - 1); i >= 0; i-- {
		h.down(i)
	}
}

꼭 필요하지는 않은, 부모, 자식 인덱스 계산기

그림을 보면 아래와 같이 각 인덱스의 부모, 왼쪽 자식, 오른쪽 자식의 인덱스 값을 계산할 수 있다.

func parentOf(i int) int {
	return (i - 1) / 2
}

func leftOf(i int) int {
	return 2*i + 1
}

func rightOf(i int) int {
	return 2*i + 2
}

 

참고링크

- YouTube: https://youtu.be/3DYIgTC4T1o

- Link: https://www.geeksforgeeks.org/building-heap-from-array/

 

Go 에서는 container/heap 패키지를 제공하지만 여기서는 그냥 날로 구현하였다. 다음에는 이 패키지를 이용해서 구현해보겠다.

- 패키지 링크: https://golang.org/pkg/container/heap/