Python heapq 实战：用小顶堆实现稳定的 Top K 排行

1. 只需要前十名，为什么要排序所有数据？

日志分析、离线报表和排行榜中，经常需要从一批记录里选出分数最高的 K 条。最直接的写法是按分数降序排序，再取前 K 条。数据已经在内存里、规模不大时，这通常也是最容易维护的方案。

但如果输入是一个只能遍历一次的迭代器，且 K 远小于记录总数 N，就可以用 Python 标准库 heapq 维护一个最多 K 个元素的小顶堆。每读入一条记录，只判断它有没有资格进入当前榜单，不必保留所有历史记录。

本文额外约定：分数相同时，先输入的记录排在前面；恰好卡在第 K 名的并列记录，也优先保留先输入者。 这个规则需要在选择阶段就体现，不能等到最后排序时再补救。

2. 堆顶不是冠军，而是当前榜单的末位

heapq 默认提供小顶堆，最小元素放在 heap[0]。在“选最大的 K 条”这个问题里，堆顶应当代表榜单中最容易被淘汰的一条：

1. 堆未满时，新记录直接入堆。
2. 堆已满时，只在新记录优于堆顶时替换堆顶。
3. 输入结束后，再对最多 K 条候选记录排序，输出最终名次。

为了同时处理分数和先后顺序，每个堆元素使用 (score, -index, record)。分数越大越好；分数相同时，较早输入者的 index 更小，-index 更大，也就更好。于是同分情况下，较晚进入的记录会更接近堆顶，优先被淘汰。

index 来自 enumerate，每条记录都不同，因此 Python 比较堆元素时不需要比较最后的字典。这也避免了直接把 (score, record) 放入堆时，同分字典无法比较而抛出 TypeError 的问题。

3. 一份可以直接运行的实现

下面约定 k 为整数，每条记录有整数类型的 score 字段；缺字段或格式错误应在数据入口处理。函数不会复制记录字典，返回的是原记录引用。

import heapq


def top_k(records, k):
    if k < 0:
        raise ValueError("k must be non-negative")
    if k == 0:
        return []

    heap = []
    for index, record in enumerate(records):
        entry = (record["score"], -index, record)
        if len(heap) < k:
            heapq.heappush(heap, entry)
        elif entry[:2] > heap[0][:2]:
            heapq.heapreplace(heap, entry)

    ranked = sorted(heap, key=lambda entry: entry[:2], reverse=True)
    return [entry[2] for entry in ranked]


if __name__ == "__main__":
    rows = [
        {"id": "A", "score": 90},
        {"id": "B", "score": 100},
        {"id": "C", "score": 100},
        {"id": "D", "score": 95},
        {"id": "E", "score": 100},
    ]
    print([row["id"] for row in top_k(iter(rows), 3)])

输出为 ['B', 'C', 'E']。如果 K 改成 2，结果为 ['B', 'C']，较晚输入但同为 100 分的 E 不会挤掉 B 或 C。

注意，heapreplace 无论新元素好坏都会替换堆顶，所以必须先做比较。另一个常见误区是直接把 heap 当作有序榜单：堆只保证父节点不大于子节点，并不保证整个列表已经排序，最后的 sorted 不能省略。

4. 用完整排序作为正确性参照

Python 的排序是稳定的，按分数降序排序后，同分记录仍保持输入顺序。因此可以用下面的简单实现作为测试参照，而不是只检查一组看起来正确的输出：

expected = sorted(rows, key=lambda row: row["score"], reverse=True)[:3]
assert top_k(iter(rows), 3) == expected

本文配套本地检查覆盖了空输入、K 为 0、K 大于记录数、K 为 1、负 K、同分截断、负分数及单次遍历输入，并使用固定随机种子生成数据，与完整排序结果逐项对照。这里验证的是功能与并列顺序，没有做生产规模的性能压测，也不据此声称固定倍数的提速。

5. 时间和空间边界

设实际保留的候选数为 M，M = min(N, K)。对于 K 大于 0 的情况，扫描阶段时间上界可写作 O(N log(K + 1))，最后排序 O(M log(M + 1))；算法额外保存 O(M) 个堆元素。K 为 0 时直接返回，不消耗输入迭代器。

这些空间结论只针对 Top K 函数自身。如果调用前已经把百万条数据读成列表，那部分内存并不会因为换成堆而自动消失。要降低端到端内存占用，输入读取也要保持增量；同时每条记录自身的大小仍会影响实际占用。

还需要明确几个业务边界：

• 这里是“选 K 条记录”，不是“取前 K 个不同分数”，也不自动按 id 去重。
• 不要在函数执行期间修改已进入堆的记录分数，否则堆里缓存的比较分数和返回字典可能不一致。
• 若改用浮点分数，应在入口明确拒绝 NaN 等不能形成正常排序关系的值；本文只处理整数分数。
• 输入必须有限，或者由调用方给出明确的批次边界；无限数据流不会自动到达最终返回点。
• K 接近 N 时，完整排序往往更简单，实际速度也需要实测；没有必要为了使用堆而使用堆。

6. 实际项目该怎么选？

如果只需要常规 Top K，优先考虑 heapq.nlargest(k, records, key=lambda row: row["score"])，标准库已经提供了简洁入口。上面的手工实现主要用来解释堆顶含义、并列规则和流式保留候选的过程，也方便以后按业务需要扩展复合比较键。

Top K 的关键不只是把排序换成堆，而是先讲清楚“什么算更好”：分数、并列先后、重复记录，以及输入何时结束。把这些规则固定下来，再用稳定排序做参照验证，算法才真正符合业务需求。