views:

1521

answers:

10

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?

+5  A: 

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
Claudiu
The try/except here breaks the iterator protocol by muffling the StopIteration, doesn't it?
David Eyk
@David Eyk: it's OK, because returning from a generator raises StopIteration anyway. The try statement in this case is superfluous.
efotinis
@efotinis: I did not know this. Thanks!
David Eyk
+18  A: 

A generator will solve your problem nicely.

def imerge(a, b):
    for i, j in itertools.izip(a,b):
        yield i
        yield j
Pramod
You should add a disclaimer - this will only work if list a is finite.
Claudiu
import itertoolsdef imerge(a, b): for i, j in zip(a,b): yield i yield jc = itertools.count(1)items = ['foo', 'bar']for i in imerge(c, items): print iI'm trying this, and this still works. len(zipped list) = min(l1, l2). So the finite length restriction is not required.
Pramod
Claudiu is correct. Try zipping two infinite generators--you will run out of memory eventually. I would prefer using itertools.izip instead of zip. Then you build the zip as you go, instead of all at once. You still have to watch out for infinite loops, but hey.
David Eyk
It will still only work if one of the arguments is a finite iterable. If they are both infinite, zip() won't work. Use itertools.izip() instead.
Thomas Wouters
Hm. I should say, Claudiu is partially correct. If either a *or* b is finite, this will work. The problem only enters in when both a *and* b are infinite, and you get an infinite loop. This is just something we have to watch out for when using infinite generators.
David Eyk
This would be my accepted answer, as I like the concise clarity of it. However, I balk at the use of zip instead of itertools.izip.
David Eyk
Right. I concede that its not ideal for 2 infinite length lists. I thought the question was about only a being finite. I'm changing it to using itertools.izip.
Pramod
Thanks. :) The question indeed specified that a was finite.
David Eyk
If you're writing any generator that calls zip() and iterates over its result, the odds are pretty good that you really want to be calling itertools.izip(). Whether or not the iterables you're zipping are finite.
Robert Rossney
In Python 3.0 zip() behaves like itertools.izip().
J.F. Sebastian
+3  A: 

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))])
Claudiu
Why do you have a restriction on the size of the first list?
Pramod
It doesn't need to be so complicated, though: `merged = chain.from_iterable(izip(items, count(1)))` will do it.
intuited
+8  A: 

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.

David Locke
Quite correct, David. @Pramod changed his answer to use izip before I noticed yours, but thanks!
David Eyk
A: 

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.

Andrea Ambu
I prefer an iterator, as I'm reading values from files of arbitrary size. I'm sure there are cases where zip is superior.
David Eyk
+1  A: 

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
John Fouhy
I always forget about enumerate! What a useful little tool, though it won't work in my particular application. Thanks!
David Eyk
+2  A: 

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.

Tom Swirly
A: 

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: 

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

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