Top K

常用综合题型

说明

维护大小为 k 的最小堆以保留最大 k 个元素；时间 O(n log k)。

分析

扫描元素时，堆中只保留目前最大的 k 个值，堆顶是这些值中最小的一个。新值若不大于堆顶即可跳过，否则替换堆顶。

图解推演

补充 k=2、输入 [4,1,7,3]。先保留 [1,4] 的最小堆；7 大于堆顶 1，替换后保留 4、7；3 不大于堆顶 4，直接丢弃。堆顶表示入选门槛，不是全局最小值；k≤0 必须单独处理，k 大于输入数时最多返回全部输入。

查看图表代码

适用前提与边界

示例返回的是堆内部顺序，不保证有序；需要按大小展示时再排序输出。

示例输入与调用

fmt.Println(TopK([]int{4, 1, 8, 3, 9}, 2))

运行输出

[8 9]

完整 Go 示例

package main

import (
	"fmt"
	"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 any) {
	*h = append(*h, x.(int))
}
func (h *IntHeap) Pop() any {
	n := len(*h) - 1
	x := (*h)[n]
	*h = (*h)[:n]
	return x
}
func TopK(a []int, k int) []int {
	h := &IntHeap{}
	for _, x := range a {
		if h.Len() < k {
			heap.Push(h, x)
		} else if k > 0 && x > (*h)[0] {
			heap.Pop(h)
			heap.Push(h, x)
		}
	}
	return append([]int(nil), (*h)...)
}

func main() {
	fmt.Println(TopK([]int{4, 1, 8, 3, 9}, 2))
}