Go 根据标准库中heap实现堆---topK问题模板

作者:guoxj
浏览:536

 剑指 Offer 40. 最小的k个数

输入整数数组 arr ,找出其中最小的 k 个数。例如,输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。

Go的实现需要注意Push和Pop方法是指针接收者,并且传入参数必须为interface,不然会失败

func getLeastNumbers(arr []int, k int) (res []int) {
    h := make(hp, len(arr))
    for idx, num := range arr {
        h[idx] = num
    }
    heap.Init(&h)
    for i := 0; i < k; i++ {
        res = append(res, (heap.Pop(&h)).(int))
    }
    return
}

type hp []int

func (h hp) Len() int { return len(h) }

// 此处为顶小堆,改为 > 后成为顶大堆
func (h hp) Less(i, j int) bool { return h[i] < h[j] }

func (h hp) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h *hp) Push(num interface{}) { *h = append(*h, num.(int)) }

func (h *hp) Pop() (res interface{}) {
    *h, res = (*h)[:h.Len()-1], (*h)[h.Len()-1]
    return
}

 




登录后回复

共有2条评论


guoxj

fwfwe

2022-05-06 22:33 回复

huangqi

sadasdasda

2022-05-06 22:33 回复