heapq.nlargest/nsmallest 时间复杂度 O(n log K)、空间 O(K),优于全排序;支持 key、生成器;不保证稳定性,需索引辅助;K 超长时自动返回全部,建议防御性截断。

用 heapq.nlargest 和 heapq.nsmallest 直接拿结果
对超大列表(比如千万级元素)找前 K 个最大/最小值,别写 sorted(lst)[-K:] 或 lst.sort() —— 它们会全量排序,时间复杂度 O(n log n),内存也扛不住。Python 标准库的 heapq.nlargest 和 heapq.nsmallest 是专为此设计的:内部用大小为 K 的堆维护候选集,时间复杂度 O(n log K),空间仅 O(K)。
实操建议:
- 当 K 远小于 n(比如 K=100,n=10⁷)时,性能优势极明显;K 接近 n 时,反而不如
sorted(因为堆操作常数更大) - 支持 key 参数,例如
heapq.nlargest(3, data, key=lambda x: x['score']) - 返回的是新列表,不修改原列表;如果原列表是生成器或迭代器,也能直接传入(避免一次性加载全部)
手动建堆比调用 nlargest 更快?别这么干
有人想“优化”:先 heapify 全列表,再反复 heappop K 次。这看似省事,但错在两点:一是 heapify 本身要 O(n) 时间且需全量加载;二是 pop K 次是 O(K log n),总复杂度退化成 O(n + K log n),比 nlargest 的 O(n log K) 更差(尤其当 K 小、n 大时)。
更关键的是:nlargest 内部根本不是先建全堆——它用的是“保留 K 个元素的可变堆”策略:遍历一次,只对当前 top-K 候选做堆调整。这是经过 C 实现高度优化的路径,自己手写 Python 循环+堆操作基本不可能更快。
立即学习“Python免费学习笔记(深入)”;
遇到重复值或自定义对象时怎么保证稳定?
heapq.nlargest 本身不保证稳定性(即相等元素的相对顺序可能改变),因为底层堆操作不记录原始索引。如果你需要稳定输出(比如按出现顺序取前 K 个最大值),得加索引辅助:
import heapq data = [3, 1, 4, 1, 5, 9] # 取前 3 个最大值,相等时按原位置靠前优先 indexed = [(x, i) for i, x in enumerate(data)] result = heapq.nlargest(3, indexed, key=lambda pair: (pair[0], -pair[1])) values = [x for x, _ in result] # [9, 5, 4]
注意点:
- key 中用
-pair[1]是为了让索引小的排前面(默认最大堆,负号翻转顺序) - 如果数据本身含不可比较对象(如字典),必须显式提供
key,否则抛TypeError: unorderable types - 对字符串、数字等内置类型,
nlargest能直接比;但自定义类没实现__lt__就会失败
内存吃紧时,流式处理大文件或数据库结果
如果列表根本装不下内存(比如从 CSV 文件读千万行),别试图先 list(reader)。直接把文件句柄或生成器传给 nlargest:
import heapq
import csv
<p>def rows_from_csv(filename):
with open(filename) as f:
reader = csv.DictReader(f)
for row in reader:
yield float(row['value'])</p><p>top_k = heapq.nlargest(100, rows_from_csv('huge.csv'))
这样全程只驻留最多 K 个元素 + 当前行,内存占用恒定。但注意:生成器只能遍历一次,如果后续还要用原数据,得另存或重读。
真正容易被忽略的是:K 的取值边界。当 K > len(iterable) 时,nlargest 会安静地返回全部元素,不会报错也不会截断——这在动态计算 K 的场景里可能掩盖逻辑错误。建议加一层 min(K, len(candidate)) 防御,尤其当 candidate 来自不确定长度的查询结果时。


















