如何合并两个Python迭代器？

Question

我有两个迭代器，一个是 list，另一个是 itertools.count 对象（即一个无限值生成器）。我想将这两个合并成一个结果迭代器，该迭代器会在两者之间交替产生值：

>>> import itertools
>>> c = itertools.count(1)
>>> items = ['foo', 'bar']
>>> merged = imerge(items, c)  # the mythical "imerge"
>>> merged.next()
'foo'
>>> merged.next()
1
>>> merged.next()
'bar'
>>> merged.next()
2
>>> merged.next()
Traceback (most recent call last):
    ...
StopIteration

最简单、最简洁的方法是什么？

- David Eyk

---

不要使用这个：list((yield next(c)) or i for i in items) - Chris_Rands

这不是 OP 寻找的内容，但在谷歌搜索“merge iterators python”时，它是第一个结果，所以我想评论一下：如果你正在寻找一种类似于归并排序的函数，将两个已排序的迭代器合并为一个更长的已排序迭代器，请使用 heapq.merge。 - Dennis

Answer 1

一个生成器可以很好地解决你的问题。

def imerge(a, b):
    for i, j in itertools.izip(a,b):
        yield i
        yield j

- Pramod

---

如果列表 a 是有限的，那么您应该添加一个免责声明 - 只有这种情况下才能起作用。 - Claudiu

Claudiu是正确的。尝试压缩两个无限生成器--你最终会耗尽内存。我更喜欢使用itertools.izip而不是zip。然后你可以一步一步地构建zip，而不是一次性全部构建。你仍然需要注意无限循环，但是嘿。 - David Eyk

它仍然只能在参数之一是有限可迭代对象时才能工作。如果它们都是无限的，zip() 将无法工作。请改用 itertools.izip()。 - Thomas Wouters

在Python 3.0中，zip()的行为类似于itertools.izip()。 - jfs

有人能为像我这样的新手澄清一下吗？如果我们使用izip，我们将能够处理从两个无限生成器中读取有限数量的元素，是吗？例如，这是izip存在的主要原因，对吗？ - Steven Lu

Answer 2

您可以做几乎与@Pramod最初建议的事情几乎完全相同。

def izipmerge(a, b):
  for i, j in itertools.izip(a,b):
    yield i
    yield j

这种方法的优点是，如果a和b都是无限的，你不会耗尽内存。

- David Locke

---

非常正确，David。@Pramod在我注意到你的回答之前更改了他的答案以使用izip，但还是谢谢！ - David Eyk

Answer 3

我也认为不需要使用itertools。

但为什么要止步于2呢？

  def tmerge(*iterators):
    for values in zip(*iterators):
      for value in values:
        yield value

可以处理从零开始的任意数量的迭代器。

更新：哎呀！评论员指出，除非所有迭代器的长度都相同，否则此代码将无法运行。

正确的代码是：

def tmerge(*iterators):
  empty = {}
  for values in itertools.zip_longest(*iterators, fillvalue=empty):
    for value in values:
      if value is not empty:
        yield value

是的，我刚刚尝试了不等长的列表以及包含{}的列表。

- Tom Swirly

---

这会用尽每个迭代器吗？我认为zip将截断到最短的那一个。我正在寻找一种合并方法，它会依次从每个迭代器中获取一个，直到它们中的每一个都被用尽。 - Thomas Vander Stichele

多尴尬啊，你说得完全正确！请看我改进后的代码。 - Tom Swirly

不用感到尴尬，你的回复和快速响应为我节省了数小时的痛苦！ - Thomas Vander Stichele

对于Python3，请将izip替换为zip。 - Stefan

@Stefan：已修复，谢谢。我是在Python 3发布两天后写的！现在Python 2已经过时了。 - Tom Swirly

Answer 4

我会这样做。这样最节省时间和空间，因为你不需要将对象一起压缩。如果a和b都是无限的，这也可以工作。

def imerge(a, b):
    i1 = iter(a)
    i2 = iter(b)
    while True:
        try:
            yield i1.next()
            yield i2.next()
        except StopIteration:
            return

- Claudiu

---

这里的 try/except 通过捕获 StopIteration 中断了迭代协议，是吗？ - David Eyk

@David Eyk：没关系，因为从生成器返回会引发StopIteration。在这种情况下，try语句是多余的。 - efotinis

Answer 5

你可以使用zip或者itertools.chain。但这只有在第一个列表是有限的的情况下才有效：

merge=itertools.chain(*[iter(i) for i in zip(['foo', 'bar'], itertools.count(1))])

- Claudiu

---

你为什么对第一个列表的大小有限制？ - Pramod

不过，并不需要这么复杂：merged = chain.from_iterable(izip(items, count(1))) 就可以了。 - intuited

Answer 6

我更喜欢这种更简洁的方式：

iter = reduce(lambda x,y: itertools.chain(x,y), iters)

- vampolo

---

在运行上面的代码行之前，在Python 3中添加from functools import reduce。 - Johan

Answer 7

Python有一个不太为人知的功能，即在生成器表达式中可以有多个for子句。这对于展开嵌套列表非常有用，比如使用zip()/izip()函数得到的列表。

def imerge(*iterators):
    return (value for row in itertools.izip(*iterators) for value in row)

- Petr Viktorin

---

肯定可以工作，但我认为嵌套的生成器表达式不太易读。如果我担心性能，我会使用这种风格。 - David Eyk

这段代码非常简洁，正如Python通常所示，但是我们该如何开始看懂它的作用呢？value for row in ...后面跟着for value in row有什么效果？这不是一个嵌套的列表推导式生成器吗？它不应该以for rowvalue in row结尾吗？还是说value被遮蔽了？ - Steven Lu

@StevenLu 基本上这是两个嵌套的循环，像这样：for row in itertools.izip(*iterators): for value in row: yield value - Petr Viktorin

Answer 8

我不确定你的应用是什么，但是你可能会发现enumerate()函数更有用。

>>> items = ['foo', 'bar', 'baz']
>>> for i, item in enumerate(items):
...  print item
...  print i
... 
foo
0
bar
1
baz
2

- John Fouhy

---

我总是忘记枚举！虽然它在我的特定应用程序中无法使用，但这是一个非常有用的小工具。谢谢！ - David Eyk

Answer 9

这里有一个优雅的解决方案：

def alternate(*iterators):
    while len(iterators) > 0:
        try:
            yield next(iterators[0])
            # Move this iterator to the back of the queue
            iterators = iterators[1:] + iterators[:1]
        except StopIteration:
            # Remove this iterator from the queue completely
            iterators = iterators[1:]

为了更好的性能，可以使用实际的队列（正如David建议的那样）：

from collections import deque

def alternate(*iterators):
    queue = deque(iterators)
    while len(queue) > 0:
        iterator = queue.popleft()
        try:
            yield next(iterator)
            queue.append(iterator)
        except StopIteration:
            pass

即使一些迭代器是有限的，而另一些迭代器是无限的，它也能正常工作：

from itertools import count

for n in alternate(count(), iter(range(3)), count(100)):
    input(n)

输出：

0
0
100
1
1
101
2
2
102
3
103
4
104
5
105
6
106

此外，它会在所有迭代器都被耗尽时正确地停止。

如果您想处理非迭代器可迭代对象（如列表），可以使用

def alternate(*iterables):
    queue = deque(map(iter, iterables))
    ...

- user76284

---

一个有趣的方法。:) 有很多种方法可以做到这一点。我想知道使用旋转的 deque() 是否比在每次迭代时重建元组更有效？ - David Eyk

Answer 10

使用 izip 和 chain 进行组合：

>>> list(itertools.chain.from_iterable(itertools.izip(items, c))) # 2.6 only
['foo', 1, 'bar', 2]

>>> list(itertools.chain(*itertools.izip(items, c)))
['foo', 1, 'bar', 2]

- Coady