代码出处

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，那么说明交换过；否则，说明没有交换过
}