golang标准库之heap

94 阅读2分钟

简介

heap堆结构是常用的数据结构,可以将多个元素按序存储,按序使用。堆是一种完全二叉树,在Go语言中标准库提供了堆的功能。

demo try

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) Pop() interface{} {
	old := *h
	o := old[len(*h)-1]
	*h = old[:len(*h)-1]
	return o
}

func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h IntHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func main() {
	h := &IntHeap{2, 3, 5}
	heap.Init(h)
	heap.Push(h, 7)
	heap.Push(h, 1)
	for h.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(h))
	}
}

一探究竟

如demo中所示,heap库提供了Init、Push、Pop、Remove等方法,同时要求自定义的结构实现 Len、Less、Swap、Pop、Push等方法

heap.Init

heap Init方法,在初始化的时候,将前一半元素,倒序做了down操作。这里因为是倒序down操作,映射到二叉树上,即为从树的下层子节点到上层父节点,逐层下沉大值的节点。

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

在down方法内,将父节点和子节点进行比较,如果子节点小于父节点,则交换,然后继续比较,直到找到合适位置。后面会对该方法进一步探究

heap.Push方法

heap Push方法,将元素插入到堆中,然后执行up操作,将子节点和父节点进行比较,如果子节点小于父节点,则交换,然后继续比较,直到找到合适位置。

	h.Push(x)
	up(h, h.Len()-1)

heap.Pop方法

heap Pop方法,将堆顶元素与尾部元素互换,然后对新的堆顶元素执行down操作,调用自定义的pop方法,弹出原堆顶最小元素。

	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()

heap.down方法

heap down方法,将父节点和子节点进行比较,如果子节点小于父节点,则交换,然后继续比较,直到找到合适位置。其中,会比较子节点的左右节点,取出左右子节点中的较小节点。保证down操作之后,父节点一定小于两个子节点。

	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 = j
	}
	return i > i0

heap.up方法

heap up方法,将子节点和父节点进行比较,如果子节点大于父节点,则交换,然后继续比较,直到找到合适位置。

	for {
		i := (j - 1) / 2 // parent
		if i == j || !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		j = i
	}

总结

整体来看,Go标准库对堆的封装简洁高效,值得我们借鉴学习,也可以将堆结构应用于日常开发工作中,提高程序的性能和效率。