优先队列(基于二叉树的堆)
代码出处
Go SDK container/heap/heap.go
Interface 接口定义
type Interface interface {
sort.Interface
Push(x interface{}) // add x as element Len()
Pop() interface{} // remove and return element Len() - 1.
}
sort.Interface是自定义排序时需要实现的接口。
type Interface interface {
// Len is the number of elements in the collection.
Len() int
// Less reports whether the element with
// index i should sort before the element with index j.
Less(i, j int) bool
// Swap swaps the elements with indexes i and j.
Swap(i, j int)
}
初始化堆
func Init(h Interface) {
n := h.Len()
for i := n/2; i >= 0; i-- { // 从第一个非叶子节点开始,遍历节点时从下往上
down(h, i, n) // 从上往下调整
}
}
添加元素
func Push(h Interface, x interface{}) {
h.Push(x) // 添加元素
up(h, h.Len()-1) // 从下往上调整
}
删除元素(取堆顶)
func Pop(h Interface) interface{} {
n := h.Len() - 1
h.Swap(0, n) // 将堆顶元素与最后一个元素交换
down(h, 0, n) // 自上往下调整
return h.Pop() // 删除堆顶元素
}
根据索引来删除某个元素
func Remove(h Interface, i int) interface{} {
n := h.Len() - 1
// 如果删除的元素是最后一个元素,那么直接删除
if n != i { // 如果删除的元素不是最后一个元素,那么交换元素后调整,再删除
h.Swap(i, n) // 将删除的元素与最后一个元素交换
if !down(h, i, n) { // 如果从上往下调整后,i 处位置元素未动,那么说明需要从下往上调整
up(h, i)
}
}
return h.Pop() // 删除最后一个元素
}
修改元素值后调整堆
func Fix(h Interface, i int) {
if !down(h, i, h.Len()) { // 如果从上往下调整后,i 处位置元素未动,那么说明需要从下往上调整
up(h, i)
}
}
从下往上调整堆
func up(h Interface, j int) {
for {
i := (j - 1) / 2 // parent
if i == j || !h.Less(j, i) { // 如果 Less 函数返回 false,那么说明不满足交换条件
break
}
h.Swap(i, j) // 交换父子元素
j = i
}
}
从上往下调整堆
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 // 如果是 true,那么说明交换过;否则,说明没有交换过
}