CODEWITHSUNDEEP
pythonlist
laboon
laboon

Reputation: 113

Iterate through 2 lists in Python x number of elements at a time?

I know you can use zip or map to iterate through them an element at a time, but how would I efficiently loop through two lists 3 elements at a time? For example:

a = [1, 2, 3, 4, 5, 6]
b = [7, 8, 9, 10, 11, 12]

zip/map would allow me to order them like

[1, 7, 2, 8, 3, 9, 4, 10, 5, 11, 6, 12]

but how would I handle a case where I want

[1, 2, 3, 7, 8, 9, 4, 5 6, 10, 11, 12]

Upvotes: 1

Views: 78

Answers (2)

iGian
iGian

Reputation: 11183

You could define your zip method, for example:

def zip_slices(iterables, n = 2):
    if n < 2: n = 1
    min_len = min([len(iterable) for iterable in iterables])
    i, size = 0, min_len
    while i < size-n+1:
        tmp = []
        for iterable in iterables:
            tmp += iterable[i:i+n]
        yield tmp
        i += n

And use it to get the desired result

a = [1, 2, 3, 4, 5, 6, 7]
b = [7, 8, 9, 10, 11, 12]

[child for parent in zip_slices((a, b), 3) for child in parent]
#=> [1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12]

The above list comprehension is flattening the output from zip_slice:

list(zip_slices((a, b), 3))
#=> [[1, 2, 3, 7, 8, 9], [4, 5, 6, 10, 11, 12]]

Upvotes: 0

rdas
rdas

Reputation: 21275

Here's a solution using the itertools recipe grouper:

from itertools import zip_longest

a = [1, 2, 3, 4, 5, 6]
b = [7, 8, 9, 10, 11, 12]

def grouper(n, iterable, fillvalue=None):
    "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return zip_longest(fillvalue=fillvalue, *args)

res = []
for x, y in zip(grouper(3, a), grouper(3, b)):
    res.extend(x + y)
print(res)

Result:

[1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12]

Upvotes: 3

Related Questions