1. 程式人生 > 其它 >[Leetcode]21.最小的k個數

[Leetcode]21.最小的k個數

輸入整數陣列 arr ,找出其中最小的 k 個數。例如,輸入4、5、1、6、2、7、3、8這8個數字,則最小的4個數字是1、2、3、4。



示例 1:

輸入:arr = [3,2,1], k = 2
輸出:[1,2] 或者 [2,1]

示例 2:

輸入:arr = [0,1,2,1], k = 1
輸出:[0]

思想:實現一個大根堆,通過維護大根堆,實現升序排序來找到最小的k個數。

import (
	"container/heap"
)
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{}){
    *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
}

這裡要注意大根堆的實現方法,以及使用方式。


題目來源:https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof