github/hello-algo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🐞 fix:

Browse Source 
If k is greater than the array itself, directly return itself
This commit is contained in:
xiaok29 2024-06-20 15:50:04 +08:00
parent a7c241609f
commit 974a75f1a0
2 changed files with 15 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
codes/go/chapter_heap/heap_test.go
View File

@ -98,4 +98,9 @@ func TestTopKHeap(t *testing.T) {
res := topKHeap(nums, k) res := topKHeap(nums, k)
fmt.Printf("最大的 " + strconv.Itoa(k) + " 个元素为") fmt.Printf("最大的 " + strconv.Itoa(k) + " 个元素为")
PrintHeap(*res) PrintHeap(*res)
k = 10
res = topKHeap(nums, k)
fmt.Printf("最大的 " + strconv.Itoa(k) + " 个元素为")
PrintHeap(*res)
} }

10
codes/go/chapter_heap/top_k.go
View File

@ -35,6 +35,16 @@ func topKHeap(nums []int, k int) *minHeap {
// 初始化小顶堆 // 初始化小顶堆
h := &minHeap{} h := &minHeap{}
heap.Init(h) heap.Init(h)
// 如果k大于数组本身返回本身数组
if len(nums) < k {
for _, num := range nums {
heap.Push(h, num)
}
return h
}
// 将数组的前 k 个元素入堆 // 将数组的前 k 个元素入堆
for i := 0; i < k; i++ { for i := 0; i < k; i++ {
heap.Push(h, nums[i]) heap.Push(h, nums[i])