题目描述

设计一个算法，找出数组中最小的k个数。以任意顺序返回这 k 个数均可。

示例：

输入： arr = [1,3,5,7,2,4,6,8], k = 4
输出： [1,2,3,4]

提示：

• 0 <= len(arr) <= 100000
• 0 <= k <= min(100000, len(arr))

思路

非常适合使用大根堆来解决：维护一个大小为 k 的大根堆，当所有元素遍历完毕后，大根堆中剩余的就是数组中最小的 k 个元素了，依次加入结果数组即可。

因为本题可以以任意顺序返回，所以虽然按照大根堆依次弹出是递减的排列，但是仍然是合法结果。如果想要得到递增排列的结果，从后往前插入到结果数组即可。

代码

type myHeap []int
func (h myHeap) Len() int {
    return len(h)
}
func (h myHeap) Less(i, j int) bool {
    return h[i] > h[j]
}
func (h myHeap) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
}
func (h *myHeap) Push(x any) {
    *h = append(*h, x.(int))
}
func (h *myHeap) Pop() any {
    x := (*h)[h.Len()-1]
    *h = (*h)[:h.Len()-1]
    return x
}

func smallestK(arr []int, k int) []int {
    var h myHeap
    var ans []int
    h = myHeap{}
    ans = make([]int, k)
    heap.Init(&h)

    for i := 0; i < len(arr); i++ {
        heap.Push(&h, arr[i])
        if h.Len() > k {
            heap.Pop(&h)
        }
    }
    
    start := 0
    for h.Len() > 0 {
        ans[start] = heap.Pop(&h).(int)  
        start++
    }

    return ans
}