Go 调用最大/小堆解决TopK问题

简介

TopK问题：在一个数组内，找到前K个最大或最小的数。

例子：剑指 Offer 40. 最小的k个数

比较简单的常用解法是排序、局部排序，除此之外，还可以使用最大/小堆。最大堆用来解决前K小的问题，最小堆用来解决前K大的问题。

堆的介绍

堆的性质如下。堆的逻辑结构为完全二叉树、底层数据结构通常为数组。对于最大堆，该二叉树父节点值皆比子节点值大，最小堆则反之。这种大小关系可以被称为堆序性。因此，对于最大堆，根为堆中最大的数。因此，最大/小堆也被称为大/小根堆。

对于堆，主要有两种操作，Push()插入新的元素，Pop()弹出堆顶（即二叉树根）元素，如果是最大堆，则弹出最大的元素，最小堆则相反。堆在进行插入和弹出操作后，都会自动调整元素位置，保证堆序性质。Go中并没有内建可以直接调用的堆容器，需要实现一些接口才可使用。

解法

以下段落假设TopK问题求前K小的数。

解决TopK前K小的数问题的思想，将数组前K个元素构建为一个最大堆，该堆最后会被作为结果集。随后从第K+1个元素扫描到末尾。因为树根为最大的数字，所以当扫描到某个数字时，该数字比树根还要小，说明有必要加入结果集。此时需要弹出堆顶，然后插入新数字。

参考"container/heap"内的定义，你需要实现sort.Interface内的方法（即Less()、Len()、Swap()），和Push()、Pop()方法。才可以创建一个堆。

type Interface interface {
    sort.Interface
    Push(x interface{}) // add x as element Len()
    Pop() interface{}   // remove and return element Len() - 1.
}

参考go的源码目录内的最小堆样例/usr/local/go/src/container/heap/example_intheap_test.go。

堆使用的切片存储。可以看到用户定义的Pop()方法只需要返回切片末尾元素并缩短一个单位，Push()方法只需要往堆使用的切片添加一个新元素。具体堆的建立和排序，都由"container/heap"包内的其他方法实现，具体他们是如何实现的，可以看源码了解。

以下是官方示例代码。

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
}

对于问题“剑指offer 40”，上面的例子里我们已经知道最小堆怎么实现了，那最大堆怎么办呢。主要区别于Less()方法，我们可以简单地定义Less()方法为，若元素i比元素j大，则视为元素i比元素j小，实现逆向排序。如果想方便复用这个堆，可以在堆的结构体内增加是最大堆还是最小堆的字段，Less()方法内根据该字段判断该如何返回。扫描一遍传入的数组，将前k个元素初始化成最大堆，从第k+1个元素开始扫描，比堆顶大，则弹出堆顶并插入新的元素。堆顶元素可以简单地通过h[0]获得。

import "container/heap"
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
// 为了实现最大堆，Less在大于时返回小于
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{}) {
    *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
}
// 最大堆
func getLeastNumbers(arr []int, k int) []int {
    h := make(IntHeap,k)
    hp:=&h
    copy(h ,IntHeap(arr[:k+1]))
    heap.Init(hp)
    for i:=k;i<len(arr);i++{
        if arr[i]<h[0]{
            heap.Pop(hp)
            heap.Push(hp,arr[i])
        }
    }
    return h
}

简易写法，由于我们需要实现sort.Interface的方法，所以我们可以直接使用sort.IntSlice结构体，该结构体实现了这个接口，无需我们再实现一次。

import (
    "container/heap"
    "sort"
)

// 继承sort.Interface的方法
type IntHeap struct {
    sort.IntSlice
}

//因为最大堆，所以覆盖Less方法，返回较大值
func (h IntHeap) Less(i, j int) bool {
    return h.IntSlice[i] > h.IntSlice[j]
}

func (h *IntHeap) Push(x interface{}) {
    h.IntSlice = append(h.IntSlice, x.(int))
}
func (h *IntHeap) Pop() interface{} {
    x := h.IntSlice[len(h.IntSlice)-1]
    h.IntSlice = h.IntSlice[:len(h.IntSlice)-1]
    return x
}

func getLeastNumbers(arr []int, k int) []int {
    if k==0{
        return []int{}
    }
    heapArr := make([]int, k)
    copy(heapArr, arr[:k])
    // 重要，取指针
    h := &IntHeap{IntSlice: heapArr}
    heap.Init(h)
    for i := k; i < len(arr); i++ {
        if x := arr[i]; x < h.IntSlice[0] {
            heap.Pop(h)
            heap.Push(h, x)
        }
    }
    return h.IntSlice

}

注意