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

Posted 2021-09-13 yizdu

TopK问题：在一个数组内，找到前K个最大或最小的数。
比较简单的常用解法是排序、局部排序，除此之外，还可以使用大/小根堆。大根堆用来解决前K小的问题，小根堆用来解决前K大的问题。

解决TopK问题的思想，将数组前K个元素构建为一个堆，随后从第K+1个元素扫描到末尾。对于前K大问题，使用小根堆，扫描过程中，如果元素大于堆顶元素，则弹出堆顶元素并插入新的元素。前K小问题则相反，使用大根堆，元素小于堆顶则弹出堆顶并插入新元素。

对于堆，主要有两种操作，Push()插入新的元素，Pop()弹出堆顶元素，如果是大根堆，则弹出最大的元素，小根堆则相反。Go中并没有内建可以直接调用的堆容器，需要实现一些接口才可使用。

参考"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. 最小的k个数
这题求前K小个数，需要使用大根堆。上面的例子里我们已经知道小根堆怎么实现了，那大根堆怎么办呢。主要区别于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
}