CODEWITHSUNDEEP
pythonpython-3.xpython-itertools
Bitto
Bitto

Reputation: 8225

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.

Upvotes: 7

Views: 743

Answers (3)

Subham
Subham

Reputation: 411

Using only list comprehension. Explanation in comments. Adding only for variation.

A = [1, 2, 3, 4, 5, 6]
B = ["A", "B", "C"]

# zipping equal length 
flattened_zip_list = [item for sublist in zip(A, B) for item in sublist]

print(flattened_zip_list )

# leftover from longest list 
extra_list = [A[index] for index in range(min(len(A), len(B)), max(len(A), len(B)) if len(A) > len(B) else B[index])]

print(extra_list)

# final list 
flattened_zip_list.extend(extra_list)

print(flattened_zip_list)

Outputs

[1, 'A', 2, 'B', 3, 'C'] # after zipping same lengths 
[4, 5, 6] # leftover values 
[1, 'A', 2, 'B', 3, 'C', 4, 5, 6] # [1, 'a', 2, 'b', 3, 'c', 4, 'd', 5, 6], the one mentioned in question 

[Program finished]

Upvotes: 0

Martijn Pieters
Martijn Pieters

Reputation: 1122182

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

Upvotes: 9

Learning is a mess
Learning is a mess

Reputation: 8277

If you want a modified version of your code, building a generator from the start (no storing list l):

import itertools
l1=[1,2,3,4,5,6]
l2=['a','b','c','d']

def flat_zip(l1,l2):
    for x,y in itertools.zip_longest(l1,l2):
        if x:
            yield x
        if y:
            yield y
it=flat_zip(l1,l2)

Though I advise for using the builtin solutions above.

Upvotes: 2

Related Questions