Golang源码 container 系列三 heap堆排序

参考
GO语言heap剖析及利用heap实现优先级队列
Golang: 详解container/heap

在Golang sort介绍了sort接口，实现自定义排序：

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

sort 包 在内部实现了四种基本的排序算法：插入排序（insertionSort）、归并排序（symMerge）、堆排序（heapSort）和快速排序（quickSort）； sort 包会依据实际数据自动选择最优的排序算法。

在container包里，专门有一个heap，实现了堆排序：

// Note that Push and Pop in this interface are for package heap's
// implementation to call. To add and remove things from the heap,
// use heap.Push and heap.Pop.
type Interface interface {
    sort.Interface
    Push(x interface{}) // add x as element Len()
    Pop() interface{}   // remove and return element Len() - 1.
}

除了sort接口，还要额外实现一下push,pop

一、官方例子：

// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This example demonstrates an integer heap built using the heap interface.
package heap_test

import (
    "container/heap"
    "fmt"
)

// An IntHeap is a min-heap of ints.
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) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
    // Push and Pop use pointer receivers because they modify the slice's length,
    // not just its contents.
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

// This example inserts several ints into an IntHeap, checks the minimum,
// and removes them in order of priority.
func Example_intHeap() {
    h := &IntHeap{2, 1, 5}
    heap.Init(h)
    heap.Push(h, 3)
    fmt.Printf("minimum: %d\n", (*h)[0])
    for h.Len() > 0 {
        fmt.Printf("%d ", heap.Pop(h))
    }
    // Output:
    // minimum: 1
    // 1 2 3 5
}

二、源码

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

其实init之后，就排好序了。关于堆排序算法，可以参考排序算法(五) 堆排序(选择排序的进化)，在heap.go中也能看到类似的逻辑：

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
}

down函数的功能非常简单：给定类型，需要down（下沉）的元素在数组中的索引，heap的长度，将该元素下沉到该元素对应的子树合适的位置，从而满足该子树为最小堆的要求。

还记得前面的那张顺序数组表示堆的图吗？结合down函数的实现：任选一个元素 i ，将其与它的子节点 2i+1 和 2i+2比较，如果元素 i 比它的子节点小，则将元素 i 与两个子节点中较小的节点交换（j），从而保证满足最小树的要求（第一次down）；子节点 j 可能也有它的子节点，继续比较、交换，直到数组末尾，或者元素 i 比它的两个子节点都小，跳出循环（）。

为什么元素 i 比它的两个子节点都小，就可以跳出循环，不再继续下去呢？这是由于，在Init函数中，第一个开始down（下沉）的元素是第 n/2 - 1 个，可以保证总是从最后一棵子树开始down（如前图，n=8或者n=9, n/2-1总是为4），因此可以保证Init->down时，如果元素 i 比它的两个子节点都小，那么该元素对应的子树，就是最小堆。

Init在遍历完毕后，可以保证，待Init的数组是一个最小堆。