Python 源代码里的算法——如何合并多个有序列表并使得结果依然有序？

摄影：产品经理

朝闻道，晚上喝酒

去年的一篇文章《一日一技：在 Python 里面如何合并多个有序列表并使得结果依然有序？》，我很自不量力地提到了“多个有序列表”。但实际上，那篇文章仅仅是合并两个有序列表而已。真正要合并多个有序列表并使结果依然有序，会难得多。

我有 A、B、C、D、E共5个有序列表，如果仅仅使用去年那篇文章的方法，那么我们需要先把 AB 合并得到列表 X，然后把 X 与 C 合并得到列表 Y，然后把 Y 与 D 合并得到列表 Z，最后把 Z 与 E 合并得到最终结果。在这个过程中，原来属于 A、B 列表的元素会被遍历4次，属于 C 列表的元素会被遍历3次……

显然，生硬套用原来的方案，效率会非常低。

那么我们如果两两组合会怎么样呢？先把 A、B 列表的元素合并，得到 X；再把 C、D列表的元素合并得到 Y、然后 XY 合并得到 Z；最后把 Z 与 E 合并得到最终结果。这个过程中，原来属于 A、B 列表的元素被遍历了3次，属于 C、D列表的元素也被遍历了3次，属于 E 列表的元素被遍历了1次。似乎好那么一丢丢，但性能依然非常差。

有什么办法能够让每个列表都只遍历一次呢？

要解决这个问题，就要用到我们的另一篇文章：一日一技：在Python里面如何获取列表的最大n个元素或最小n个元素？中涉及到的一个数据结构—最小堆（又叫小顶堆）。

最小堆本质是一个二叉树，并且父节点总是小于等于子节点。根节点总是最小的。

这个算法需求，Python 已经帮我们做好了，那就是heapq.merge()，相关代码如下：

import heapq


A = [1, 2, 3, 4, 5, 6]
B = [-6, -5, -4, -3, -2, -1]
C = [-100, -20, 0, 12, 15, 18]
D = [-1, 13, 56, 100]
E = [3, 7]
result = [x for x in heapq.merge(A, B, C, D, E)]
print(result)

运行结果如下：

那么，大家想不想知道，这背后的原理是什么？实际上，这个原理说起来很简单：

现在，我们分别从 ABCDE 三个有序列表中，取出最小的元素（下标为0的元素），并把他们构成一个最小堆。

然后从最小堆里面取出堆顶元素，放到结果列表中。接下来，我们从刚才取出的这个元素原来所在的列表中，再取一个元素出来，放入最小堆中。如果它依然是最小的，那么它直接就在堆顶；如果它不是堆中最小的，那么堆顶会变成另一个元素。把堆顶元素取出来，放入结果列表中。接下来从这个被取出来的堆顶元素原来所在的列表中，取最小的元素，继续放入堆中……

一开始有5个列表，所以堆中始终保持5个元素。后来有一个列表空了，那么此时堆中始终保持4个元素……最后直到只剩1个列表时，直接拼接到结果列表末尾即可。

说起来确实简单，但是：

Talk is cheap, show me the code. ”

信不信，虽然我已经把原理告诉了你，但是现在给你30分钟的时间，你可能也无法把代码写出来。

你不要自怨自艾，虽然我想通这个原理只花了5分钟，但是我用了半个小时也没有把代码完整写出来。

所以，我们来看看 Python 的源代码，看看它是怎么写的。heapq.merge的源代码在Python 的 heapq.py 文件中。你可以在 PyCharm 中，输入：

import heapq
heapq.merge

Windows/Linux 按住 Ctrl 键，macOS 按住 Command 键，鼠标左键点击merge这个词，自动跳转到源代码中，如下图所示：

核心代码非常简单，只有16行。

图中第332行的h列表将会实现一个堆。第335行-344行，大家可以忽略，这里是根据输入的多个有序列表是从小到大还是从大到小做的针对性处理。我们解释原理的时候，假设输入的多个列表都是从小到大的有序列表。

正餐从第347行开始。

for order, it in enumerate(map(iter, iterables))

看起来这么复杂，实际上我给你简化一下你就知道了：

for order, it in enumerate([iter(A), iter(B), iter(C), iter(D), iter(E)])

所以实际上我们在这个循环里面，分别会得到如下几组 order 和 it：

order = 0, it = iter(A)
order = 1, it = iter(B)
order = 2, it = iter(C)
order = 3, it = iter(D)
order = 4, it = iter(E)

为什么需要把一个列表传给iter函数呢？这是为了一个一个地取出列表中的元素。

我们知道，当你使用列表.pop(0)弹出列表下标为0的元素时，列表后面的元素会依次向前移动1位。这就导致 列表.pop(0)时间复杂度为 O(n)。

但是，当我们使用iter(列表)把一个列表转换为迭代器以后，只需要执行迭代器.__next__()就能获取下标为0的元素，并且时间复杂度为 O(1)。

iter(列表)的工作原理，可以近似等价于下面这段代码：

def iter(A):
    for element in A:
        yield element

回到原来的代码中，第349行，把 it.__next__函数赋值给 next变量，从而可以实现next()的效果与it.__next__()相同。从这里也可以看到，这段代码有点老了，因为新版的 Python 自带了next()函数，next(it)就等价于it.__next__()。可以推测出，写这段代码的时候，Python 还没有自带的next()函数。

第350行，向 h 列表中添加一个列表：[next(), order * direction, next]，如果我们使用A = [1, 2, 3, 4, 5, 6]来作为例子的话，那么首先添加到 h 列表中的数据是：[1, 0, it_A]，其中的 it_A 是 A 取掉第一个元素以后的迭代器。（如果再次执行 it_A.__next__()）那么返回的是数字2.

代码347行-352行运行完成以后，h 列表是这样的：

h = [
[1, 0, lt_A],
[-6, 1, lt_B],
[-100, 2, lt_C],
[-1, 3, lt_D],
[3, 4, lt_E]
]

第353行代码，把列表 h 变成一个最小堆。于是列表 h 中的数据顺序发生了变化：

h = [
[-100, 2, lt_C],
[-6, 1, lt_B],
[-1, 3, lt_D],
[1, 0, lt_A],
[3, 4, lt_E]
]

这里我们停一下，大家考虑一个问题，为什么 h 中的子列表，每个都有3个元素？第一个元素是原来各个列表中最小的数字，这个很好理解，将会用来排序。但为什么代码里面有一行order * direction，并放到列表的第二项？

要解答这个问题，我们就需要知道，Python 的列表的大小对比规则。

有两个列表：a = [1, 2, 3]和b=[4, 5, 6]，在 Python 里面执行b > a会返回 True，如下图所示：

实际上，它的原理就是一项一项对比列表中的每个元素。所以当判断[4, 5, 6]是否大于[1, 2, 3]时，是首先判断4是否大于1，发现大于，于是就停止对比，直接返回 True。

如果第一个元素相同，就再对比各自的第二个元素。由于要对每个元素都进行对比，这就要求列表中的当前被对比的元素是可以比较大小的。但是迭代器是不能对比大小的。

回来我们的算法中，当 h 中的子列表第一个数字相同时，就会开始对比第二个数。由于第二个数是按 A-E 分1-5，所以此时必定可以比较出结果，于是就是实现了：取最小的，如果相同就取第一个，这样的功能。不会去对比第三项迭代器。

我们继续来看 Python 的源代码。先看第363-368行。如果我们传给heapq.merge只有1个有序列表。那么直接把里面每个元素抛出去即可。如果超过1个列表，进入354-362行，也就是核心算法部分。

第356行，首先取出堆顶元素，也就是最小的这个元素。由于这个子列表有3项，分别赋值给value/order/next。其中 value 就是最小的一个元素，直接抛出给外层。

然后获取当前这个元素原来所在列表剩下数据中最小的元素。

这里就是这个算法精妙的地方了，它通过一个列表把当前的数字和这个数字原来所在的列表的迭代器绑定在了一起。找到了当前这个数字，自然就能找到它原来属于的那个列表的迭代器。并且执行迭代器.__next__()就可以获取到下一条数据。

接下来，第359行，把堆顶列表的下标为0的项替换为新的值（原来所在列表的第二个值）。

第360行，把当前堆顶的列表替换为新的列表。个人觉得这个地方实际上不需要替换，因为在第359行修改s[0]的时候，就已经更新堆顶这个列表下标为0的值了。但是由于更新堆顶以后，堆顶可能并不是最小的，于是需要使用heapify重新调整堆。它使用heapreplace也没有什么问题，也能实现调整的目的。

调整完成以后，进行下一轮循环，继续弹出堆顶列表下标为0的元素，更新堆顶……

由于不同的列表长度不同，当某个列表耗尽以后，迭代器就会抛出StopIteration异常，于是堆元素就减少，直到减到1个以后，逐项抛出剩下的元素即可。

通过这个算法，每个列表只需要遍历1次，效率大幅度提升。

在本文中，我们使用的是列表。如果把有序列表换成有序链表，解答思路完全一样，并且还省略了转换为迭代器的一步，代码还要少一些。换成链表以后，这就是 LeetCode 上难度为 Hard（困难）的题目了。但如果你认真理解本文，你会发现，Hard 难度实际上也不过如此。