go_heap初探

2021-05-04

什么是heap

堆是一种特殊的树结构,它的根节点总是全局最小或者最大,自然也就衍生出了最大根堆和最小根堆

go的heap实现

go的heap实现了如下接口

type Interface interface {
	sort.Interface
	Push(x interface{}) // add x as element Len()
	Pop() interface{}   // remove and return element Len() - 1.
}

初始化时进行了如下操作

func Init(h Interface) {
	// heapify
	n := h.Len() 
	for i := n/2 - 1; i >= 0; i-- {
		down(h, i, n)
	}
}

// 向下调整(pop的时候需要调换根节点和尾结点)
func down(h Interface, i0, n int) bool {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
			j = j2 // = 2*i + 2  // right child  找一个更小的,维持小根堆
		}
		if !h.Less(j, i) {
			break
		}
		h.Swap(i, j) // 也就是i的位置值一直比较大
		i = j
	}
	return i > i0
}

// 向上调整,一旦一个新值添加进来的时候需要放到尾部向上调整 
func up(h Interface, j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		j = i
	}
}

// Push pushes the element x onto the heap.
// The complexity is O(log n) where n = h.Len().
func Push(h Interface, x interface{}) {
	h.Push(x)
	up(h, h.Len()-1)
}

// Pop removes and returns the minimum element (according to Less) from the heap.
// The complexity is O(log n) where n = h.Len().
// Pop is equivalent to Remove(h, 0).
func Pop(h Interface) interface{} {
	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

总结

go的heap用起来比较麻烦,因为需要自己去实现interface中定义的各个方法,来指定大根堆或者小根堆,源码中也没有默认实现,使用起来不如java的顺手。

下面是最常用的int型实现

// An IntHeap is a min-heap of ints.
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] } // 可以取相反,来做大根堆
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
	// Push and Pop use pointer receivers because they modify the slice's length,
	// not just its contents.
	*h = append(*h, x.(int))
}

// 常规pop操作
func (h *IntHeap) Pop() interface{} {
	old := *h // 定义一个新变量
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// This example inserts several ints into an IntHeap, checks the minimum,
// and removes them in order of priority.
func Example_intHeap() {
	h := &IntHeap{2, 1, 5}
	heap.Init(h)
	heap.Push(h, 3)
	fmt.Printf("minimum: %d\n", (*h)[0])
	for h.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(h))
	}
	// Output: // 这是栗子写法
	// minimum: 1
	// 1 2 3 5
}