Question

I have two iterators, a list and an itertools.count object (i.e. an infinite value generator). I would like to merge these two into a resulting iterator that will alternate yield values between the two:

>>> 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

What is the simplest, most concise way to do this?

Was it helpful?

Solution

A generator will solve your problem nicely.

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

OTHER TIPS

You can do something that is almost exaclty what @Pramod first suggested.

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

The advantage of this approach is that you won't run out of memory if both a and b are infinite.

I also agree that itertools is not needed.

But why stop at 2?

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

handles any number of iterators from 0 on upwards.

UPDATE: DOH! A commenter pointed out that this won't work unless all the iterators are the same length.

The correct code is:

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

and yes, I just tried it with lists of unequal length, and a list containing {}.

I'd do something like this. This will be most time and space efficient, since you won't have the overhead of zipping objects together. This will also work if both a and b are infinite.

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

You can use zip as well as itertools.chain. This will only work if the first list is finite:

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

I'm not sure what your application is, but you might find the enumerate() function more useful.

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

I prefer this other way which is much more concise:

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

One of the less well known features of Python is that you can have more for clauses in a generator expression. Very useful for flattening nested lists, like those you get from zip()/izip().

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

Here is an elegant solution:

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:]

Using an actual queue for better performance (as suggested by 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

It works even when some iterators are finite and others are infinite:

from itertools import count

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

Prints:

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

It also correctly stops if/when all iterators have been exhausted.

If you want to handle non-iterator iterables, like lists, you can use

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

Use izip and chain together:

>>> 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]

Why is itertools needed?

def imerge(a,b):
    for i,j in zip(a,b):
        yield i
        yield j

In this case at least one of a or b must be of finite length, cause zip will return a list, not an iterator. If you need an iterator as output then you can go for the Claudiu solution.

Using itertools.izip(), instead of zip() as in some of the other answers, will improve performance:

As "pydoc itertools.izip" shows: "Works like the zip() function but consumes less memory by returning an iterator instead of a list."

Itertools.izip will also work properly even if one of the iterators is infinite.

A concise method is to use a generator expression with itertools.cycle(). It avoids creating a long chain() of tuples.

generator = (it.next() for it in itertools.cycle([i1, i2]))
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top