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 个元素。可以结合哈希表 + 小顶堆实现:

  1. 用哈希表统计每个元素的出现频次;
  2. 维护一个大小为 K 的小顶堆,其中存储 (频次, 元素)
  3. 遍历哈希表,用频次作为比较依据,保持堆大小 K;
  4. 最后堆中剩下的即为高频 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]