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

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
}
阅读 2.1k

yizdu@Pop_OS! ~

4 声望
0 粉丝
0 条评论

yizdu@Pop_OS! ~

4 声望
0 粉丝
文章目录
宣传栏