Golang heap源码简单走读

golang heap小根堆源码走读

heap概览

在golang中，通过heap给出了一个实现小根堆的接口。

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

由于小根堆中，需要根据容器中的元素大小来进行比较以确定元素在堆中的位置。因此也需要实现sort的接口。

type Interface interface {
	Len() int
	Less(i, j int) bool
	Swap(i, j int)
}

因此，如果需要实现一个heap，需要实现Push(),Pop()（从堆顶弹出元素）,Len(),Less()（进行元素之间比较的方法)以及Swap()方法。

heap已经实现的方法

heap的初始化

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

在heap的初始化方法中，值得注意的是，不同于java方法，在调用这个方法的时候，堆中的元素容器应该已经完成了数据的存放，这里的初始化不是初始化存储空间返回一个空的容器，而是对一个已经充满元素的容器进行堆排序。首先会获取堆的大小n，取决于堆的性质，将会从堆的n/2-1位置开始通过down()方法初始化，这个位置是堆中倒数第二层最后一个拥有子节点的节点，之后将会不断从当前层不断向左，遍历完当前层之后再从上一层的最右边开始。

func down(h Interface, i0, n int) bool {
	i := i0 // i是方法中当前节点的下标
	for {
		j1 := 2*i + 1 // 获取当前节点的左子节点
		if j1 >= n || j1 < 0 { 
			break // 如果当前节点的左子节点不存在则结束这一轮
		}
		j := j1 // j为当前节点的左子节点
		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) { // 此处Less()方法为上文提到需要扩展实现的元素比较方法
			j = j2 // 将节点的左节点和右节点进行比较，j赋值为其中较小的下标
		}
		if !h.Less(j, i) {
			break 
		}
		h.Swap(i, j) // 如果当前节点小于其子节点中的较小值，通过扩展的Swap()方法交换,将较小的值替换到父节点上，并且将会从当前节点继续往下与其子节点比较，直到达到最底层或者其两个子节点逗都比自己大
		i = j
	}
	return i > i0
}

顾名思义，down()实则是将一个节点不断从堆中不断下移直到达到其对应位置的一个过程。由于n/2-1保证了遍历的开始位置是堆中最后一个拥有子节点的位置，因此可以保证最后一层每一个子节点都能参与到一轮小根堆的初始化。一轮down()下来，将会保证容器中的最小值将会处于堆顶，只需要通过Pop()方法将栈顶的元素弹出就可以得到堆中最小的值。这可以适用于优先级队列等场景的实现。

heap的添加

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

当通过Push()方法往堆中添加元素的时候，可以先简单的将元素放到最后一层的叶子节点上，之后通过up()方法从最后一个子节点开始，将该新元素从堆底up到堆中的一个合适的位置上。

func up(h Interface, j int) {
	for {
		i := (j - 1) / 2 // 当前节点的父节点
		if i == j || !h.Less(j, i) {
			break
		}
		h.Swap(i, j) // 如果当前节点小于其父节点，将该节点与其父节点交换，并继续从新的位置向其父节点进行比较，直到下标为0达到堆顶或者遇到比其更小的父节点
		j = i
	}
}

对应down()的取名,up()是将一个节点从当前位置不断往堆的更上层前进的过程。这里的Push()操作的时间复杂度为O(logn)。

heap的弹出

func Pop(h Interface) interface{} {
	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

Heap的弹出，只需要将堆顶的元素与堆的末尾元素调换，将其放到堆的末尾默认移除，不参与后续堆的调整即可。之后将当前所处的堆顶元素通过down()方法不断下移到其应该处在的位置即可。