Skip to content
Advertisement

Create iterator to return elements from each iterable one by one

I was learning the itertools module and I am trying to make an iterator to return each element from the iterables provided as input.

Agruments   Results
p, q, …     p0, q0, … plast, qlast

with one more rider that if say the lists are not of the same length then next(it) should return elements from the longer list when the shorter one runs out.

Attempt at solution

import itertools
l1=[1,2,3,4,5,6]
l2=['a','b','c','d']
l=[]
for x,y in itertools.zip_longest(l1,l2):
    l.extend([x,y])
it=iter(x for x in l if x is not None)

Which kind of solves my problem

print(list(it))

Outputs:

[1, 'a', 2, 'b', 3, 'c', 4, 'd', 5, 6]

Is there an easier or better way to do this? I searched for a solution on SO and was not able to get one.

Advertisement

Answer

You can use itertools.chain.from_iterable() to flatten the sequence, and use a generator expression to filter out the None values:

from itertools import chain, zip_longest

it = (v for v in chain.from_iterable(zip_longest(l1, l2)) if v is not None)

Rather than use None as the sentinel value, you may want to use a dedicated sentinel so you can use None in the input list:

_sentinel = object()
flattened = chain.from_iterable(zip_longest(l1, l2, fillvalue=_sentinel))
it = (v for v in flattened if v is not _sentinel)

If you want to filter out falsey values, then you can also use filter(None, ...):

it = filter(None, chain.from_iterable(zip_longest(l1, l2)))

Demo:

>>> from itertools import chain, zip_longest
>>> l1 = [1, 2, 3, 4, 5, 6]
>>> l2 = ['a', 'b', 'c', 'd']
>>> it = (v for v in chain.from_iterable(zip_longest(l1, l2)) if v is not None)
>>> list(it)
[1, 'a', 2, 'b', 3, 'c', 4, 'd', 5, 6]

and with a local sentinel:

>>> l1 = [1, None, 2, None, 3, None]
>>> l2 = ['a', 'b', 'c', 'd']
>>> _sentinel = object()
>>> flattened = chain.from_iterable(zip_longest(l1, l2, fillvalue=_sentinel))
>>> it = (v for v in flattened if v is not _sentinel)
>>> list(it)
[1, 'a', None, 'b', 2, 'c', None, 'd', 3, None]

The itertools recipes section also has:

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))
Advertisement