问题
输入整数数组 arr ,找出其中最小的 k 个数。例如,输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。
代码
注意看看,用 go
实现堆是如何实现的?
package main
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] }
// Push and Pop use pointer receivers because they modify the slice's length,
// not just its contents.
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 {
if k == 0 {
return []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
}