go源码阅读 - heap实现

go 源码阅读 - heap 实现

go 中对于堆的实现很简单，代码不到100行，位于 src/container/heap 包下。

Push

// src/container/heap/heap.go
// 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)
}

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 完之后， 会对 push 进去的元素进行 up 操作，up 方法会向上遍历，然后比较两者的大小(分最大堆和最小堆)，如果满足要交换的条件(最大堆是大于父节点，最小堆是小于父节点)，就交换两者的位置，一直到找到该元素正确的位置。时间复杂度 O(log(n))

Pop

// 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()
}

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

要 Pop 一个元素时，Pop 出来的是最大值或者最小值，执行 Pop 时先将首位元素和末位元素交换位置，那么末尾元素就是堆首的元素了(最大值/最小值)，然后堆除去末尾元素的剩下值的堆首进行 down 操作，找到堆首元素应该在的正确位置。最后调用当前对象的 Pop 方法，把末尾元素pop出来，那么pop出来的就是该堆的最大值/最小值，并且剩下的堆也是一个最大堆/最小堆了。时间复杂度 O(log(n))