Top K 问题堆排序解法
FreeGuideOnline
6阅读
2026-07-10
python import heapq
def top_k_largest(nums, k): if k <= 0: return [] min_heap = [] for num in nums: if len(min_heap) < k: heapq.heappush(min_heap, num) else: # 堆已满,只有更大的值才值得进入 if num > min_heap[0]: heapq.heapreplace(min_heap, num) return min_heap # 此时堆中即为前 K 大,堆顶是第 K 大
示例
data = [3, 1, 5, 12, 2, 11, 4, 8] print(top_k_largest(data, 3)) # 输出 [8, 11, 12](顺序不定)
**关于 `heapreplace`**:它是弹出堆顶并插入新元素的高效操作,相当于 `heappop` + `heappush` 但更省时。
#### 6.2 求最小的 K 个元素
将逻辑改为大顶堆。Python 标准库只有小顶堆,可通过存入**负值**变相构造大顶堆。
```python
def top_k_smallest(nums, k):
if k <= 0:
return []
max_heap = [] # 用负值模拟大顶堆
for num in nums:
if len(max_heap) < k:
heapq.heappush(max_heap, -num)
else:
# 如果新元素更小(负值更大),替换
if -num > max_heap[0]:
heapq.heapreplace(max_heap, -num)
# 恢复正值并返回
return [-x for x in max_heap]
# 测试
print(top_k_smallest(data, 3)) # 输出 [1, 2, 3]
7. 变体:Top K 高频元素
给定一个数组,找出出现频率最高的 K 个元素。可以结合哈希表 + 小顶堆实现:
- 用哈希表统计每个元素的出现频次;
- 维护一个大小为 K 的小顶堆,其中存储
(频次, 元素); - 遍历哈希表,用频次作为比较依据,保持堆大小 K;
- 最后堆中剩下的即为高频 Top K。
from collections import Counter
def top_k_frequent(nums, k):
freq = Counter(nums)
min_heap = []
for num, count in freq.items():
if len(min_heap) < k:
heapq.heappush(min_heap, (count, num))
else:
if count > min_heap[0][0]:
heapq.heapreplace(min_heap, (count, num))
return [num for count, num in min_heap]