如何从两个有序可迭代对象中高效获取前N个有序元素？

当然可以！既然i1和i2已经按同一键值完成排序，我们完全可以借助itertools的工具，用归并算法的思路来高效提取前N个元素，避免把整个可迭代对象加载到内存做全量排序——这在处理大数据或无限迭代器时优势特别明显。

核心思路 #

因为两个输入都是已排序的，我们不需要重新对所有元素排序，只需要像归并排序那样，每次从两个迭代器的当前头部选出最小的元素，直到收集够N个为止。虽然标准库中heapq.merge是干这件事的最佳工具，但它不属于itertools；如果必须仅用itertools实现，我们可以手动基于itertools.tee、itertools.islice等工具构建归并逻辑。

实现代码 #

1. 基于itertools的归并函数

先实现一个归并两个已排序迭代器的生成器，用itertools.tee复制迭代器来避免消费原输入：

import itertools

def merge_sorted(i1, i2, key=None):
    # 默认用元素本身作为排序键
    key_func = key if key is not None else lambda x: x
    
    # 复制迭代器，防止修改原可迭代对象
    it1, it1_rest = itertools.tee(i1)
    it2, it2_rest = itertools.tee(i2)
    
    try:
        val1 = next(it1)
        val2 = next(it2)
        while True:
            if key_func(val1) <= key_func(val2):
                yield val1
                val1 = next(it1)
            else:
                yield val2
                val2 = next(it2)
    except StopIteration:
        # 其中一个迭代器耗尽后，输出另一个的剩余元素
        yield from it1_rest
        yield from it2_rest

2. 提取前N个元素

用itertools.islice直接从归并后的迭代器中截取前N个元素，不需要加载全部数据：

def get_top_n(i1, i2, n, key=None):
    merged_iter = merge_sorted(i1, i2, key)
    return list(itertools.islice(merged_iter, n))

为什么这比全量排序更好？ #

• 内存友好：不需要把i1和i2全部转成列表，哪怕是TB级的数据流或者无限迭代器，只要N有限就能正常工作
• 速度更快：全量排序的时间复杂度是O((M+K)log(M+K))，而这种归并取前N的方式时间复杂度是O(N)（当N远小于总元素数时，效率提升非常显著）

注意点 #

• 必须保证i1和i2已经按指定的key升序排列，否则归并结果会出错
• 如果需要降序的前N个元素，可以先给输入迭代器加上反转逻辑（比如用reversed，但要注意原迭代器是否支持反转），或者调整归并时的比较逻辑

内容的提问来源于stack exchange，提问作者Roy Smith