You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
52 lines
1.3 KiB
52 lines
1.3 KiB
// File: top_k.go
|
|
// Created Time: 2023-06-24
|
|
// Author: Reanon (793584285@qq.com)
|
|
|
|
package chapter_heap
|
|
|
|
import "container/heap"
|
|
|
|
type minHeap []any
|
|
|
|
func (h *minHeap) Len() int { return len(*h) }
|
|
func (h *minHeap) Less(i, j int) bool { return (*h)[i].(int) < (*h)[j].(int) }
|
|
func (h *minHeap) Swap(i, j int) { (*h)[i], (*h)[j] = (*h)[j], (*h)[i] }
|
|
|
|
// Push heap.Interface 的方法，實現推入元素到堆積
|
|
func (h *minHeap) Push(x any) {
|
|
*h = append(*h, x.(int))
|
|
}
|
|
|
|
// Pop heap.Interface 的方法，實現彈出堆積頂元素
|
|
func (h *minHeap) Pop() any {
|
|
// 待出堆積元素存放在最後
|
|
last := (*h)[len(*h)-1]
|
|
*h = (*h)[:len(*h)-1]
|
|
return last
|
|
}
|
|
|
|
// Top 獲取堆積頂元素
|
|
func (h *minHeap) Top() any {
|
|
return (*h)[0]
|
|
}
|
|
|
|
/* 基於堆積查詢陣列中最大的 k 個元素 */
|
|
func topKHeap(nums []int, k int) *minHeap {
|
|
// 初始化小頂堆積
|
|
h := &minHeap{}
|
|
heap.Init(h)
|
|
// 將陣列的前 k 個元素入堆積
|
|
for i := 0; i < k; i++ {
|
|
heap.Push(h, nums[i])
|
|
}
|
|
// 從第 k+1 個元素開始，保持堆積的長度為 k
|
|
for i := k; i < len(nums); i++ {
|
|
// 若當前元素大於堆積頂元素，則將堆積頂元素出堆積、當前元素入堆積
|
|
if nums[i] > h.Top().(int) {
|
|
heap.Pop(h)
|
|
heap.Push(h, nums[i])
|
|
}
|
|
}
|
|
return h
|
|
}
|