golang에서 heap 사용하기

파이썬 같은 언어는 기본으로 내장되어 있는 heapq가 사용하기 상당히 쉽고 편리한데에 비해 go에서 heap을 사용하려면 꽤나 복잡한 작업을 해야합니다. 이번 글에서는 go에서 heap을 어떻게 사용하는지 공식문서예제에 있는 코드로 알아보겠습니다. 

 

공식문서의 예제 링크: golang.org/src/container/heap/example_pq_test.go

 

func Pop(h Interface) interface{}

type Interface interface {
    sort.Interface
    Push(x interface{}) 
    Pop() interface{}   
}

Interface type을 구현하기 위한 함수들은 위처럼 정의 되어 있습니다.

Len() int
Less(i, j int) bool
Swap(i, j int)

위 5개의 함수를 구현해야 합니다.

 

그럼 위 5개의 함수를 정의한 타입을 만들어 보겠습니다. (golang 공식문서 예제입니다.)

type Item struct {
    value    string 
    priority int    
    index int 
}
 
type PriorityQueue []*Item
 
func (pq PriorityQueue) Len() int { return len(pq) }
 
func (pq PriorityQueue) Less(i, j int) bool {
    return pq[i].priority > pq[j].priority  // priority의 값이 클수록 우선순위가 높음 
}
 
func (pq PriorityQueue) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
    pq[i].index = i
    pq[j].index = j
}
 
func (pq *PriorityQueue) Push(x interface{}) {
    n := len(*pq)
    item := x.(*Item)
    item.index = n
    *pq = append(*pq, item)
}
 
func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    old[n-1] = nil  // avoid memory leak
    item.index = -1 // for safety
    *pq = old[0 : n-1]
    return item
}

Item 구조체들을 가지고 있는 PriorityQueue 타입입니다. Len, Less, Swap 함수들의 설명은 생략하겠습니다.

PriorityQueue는 Item 구조체포인터의 슬라이스인데  []Item과 []*Item의 차이는 아래와 같습니다.

package main
 
import "fmt"
 
type Item struct {
    node   int
    weight int
}
 
type PriorityQueue []Item
 
func main() {
    i1 := Item{1, 4}
    i2 := Item{2, 7}
    i3 := Item{3, 19}
 
    pq := PriorityQueue{i1, i2, i3}
    pq[0].weight = 11
    fmt.Println(i1.weight)
}

[]Item으로 하는 경우 i1의 weight는 11로 바뀌지 않고 4로 유지됩니다.

package main
 
import "fmt"
 
type Item struct {
    node   int
    weight int
}
 
type PriorityQueue []*Item
 
func main() {
    i1 := &Item{1, 4}
    i2 := &Item{2, 7}
    i3 := &Item{3, 19}
 
    pq := PriorityQueue{i1, i2, i3}
    pq[0].weight = 11
    fmt.Println(i1.weight)
}

[]*Item의 경우 i1의 weight가 11로 바뀌게 됩니다.

 

 

Push함수를 보면 PriorityQueue의 포인터 타입을 받고 있습니다. (pq *PriorityQueue)
그 이유는 아래 예제를 보시면 이해할 수 있습니다.

package main
 
import "fmt"
 
type PriorityQueue []int
 
func (pq PriorityQueue) Push(x int) {
    pq = append(pq, x)
    fmt.Println(pq)
}
 
func main() {
    pq := PriorityQueue{}
    pq.Push(1)
    pq.Push(2)
    pq.Push(3)
    fmt.Println(pq)
}

포인터 타입이 아닌 (pq PriorityQueue)의 형태로 구현하게 되면 값 복사가 이루어지기 때문에 pq에 1, 2, 3 모두 들어가지 않고
빈 슬라이스가 출력되어집니다.

위 예제의 출력값은 아래와 같습니다.
[1]  // Push 함수 내에서 출력됨 (빈 슬라이스인 pq에다가 1을 넣고 출력된 것)
[2] // Push 함수 내에서 출력됨 (빈 슬라이스인 pq에다가 2를 넣고 출력된 것)
[3] // Push 함수 내에서 출력됨(빈 슬라이스인 pq에다가 3을 넣고 출력된 것)
[]   // main에서 선언한 pq (값 복사로 Push함수가 실행되어서 결국 아무 값도 들어오지 않음)

 

다시 공식문서 예제로 돌아와서 Push함수의 x.(*Item)에 관해서는 아래 링크의 인터페이스 형 변환 부분을 참고해주세요.
golang 인터페이스 활용: up-to-date-items.tistory.com/116

그 밑의 코드들은 슬라이스의 마지막에 추가해주는 것입니다. 그런데 이런 의문이 들 것입니다. 힙 자료구조의 특성상 새로운 원소가 Push 됬으면 upheap과정이 추가되어야 합니다.

func Push(h Interface, x interface{}) {
    h.Push(x)
    up(h, h.Len()-1)
}

container/heap의 Push함수를 보면 위와 같이 구현되어 있습니다. 즉, 우리가 Push함수만 구현해주면 upheap의 과정은 내부적으로 구현이 되어 있습니다.

 

이제 Pop함수를 알아보겠습니다.

Push함수와 마찬가지로 원본 슬라이스 접근을 위해 포인터 형태로 구현되어 있습니다. (pq *PriorityQueue)

그리고 슬라이스의 마지막 원소(우선순위가 제일 높은 원소)를 pop시키고 인덱스와 원본 슬라이스를 조정합니다.

Push함수처럼 downheap과정은 container/heap의 Push함수에 내부적으로 구현되어 있습니다.

 

이렇게 구현을 하였으면 아래처럼 heap.Push, heap.Pop을 이용하여 사용하면 됩니다.

pq := make(PriorityQueue, len(items))
heap.Push(&pq, item)

(단순 예제를 위한 코드라 items에 관한 코드는 생략하였습니다.)

Pusu와 Pop함수를 사용할 때는 역시 값 복사를 피하기 위해 pq의 주소를 인자로 넘겨주어야 합니다.

또한, heap.Pop함수의 리턴 값은 interface{} 타입입니다. 그러므로 아래처럼 인터페이스 형변환이 필요합니다.

item := heap.Pop(&pq).(*Item)