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