Attempting Python list comprehension with two variable of different ranges

Question

I'm trying to generate a list quickly with content from two different arrays of size n and n/2. As an example:

A = [70, 60, 50, 40, 30, 20, 10, 0]
B = [1, 2, 3, 4]

I wish to generate something like

[(A[x], B[y]) for x in range(len(A)) for y in range(len(B))]

I understand the second for statement is the nested for loop after the "x" one. I'm trying to get the contents of the new array to be

A[0], B[0]
A[1], B[1]
A[2], B[2]
A[3], B[3]
A[4], B[0]
A[5], B[1]
A[6], B[2]
A[7], B[3]

Could anyone point me in the right direction?

Martijn Pieters · Accepted Answer · 2016-02-05 02:01:30Z

21

Don't use nested loops; you are pairing up A and B, with B repeating as needed. What you need is zip() (to do the pairing), and itertools.cycle() (to repeat B):

from itertools import cycle

zip(A, cycle(B))

If B is always going to be half the size of A, you could also just double B:

zip(A, B + B)

Demo:

>>> from itertools import cycle
>>> A = [70, 60, 50, 40, 30, 20, 10, 0]
>>> B = [1, 2, 3, 4]
>>> zip(A, cycle(B))
[(70, 1), (60, 2), (50, 3), (40, 4), (30, 1), (20, 2), (10, 3), (0, 4)]
>>> zip(A, B + B)
[(70, 1), (60, 2), (50, 3), (40, 4), (30, 1), (20, 2), (10, 3), (0, 4)]

For cases where it is not known which one is the longer list, you could use min() and max() to pick which one to cycle:

zip(max((A, B), key=len), cycle(min((A, B), key=len))

or for an arbitrary number of lists to pair up, cycle them all but use itertools.islice() to limit things to the maximum length:

inputs = (A, B)  # potentially more
max_length = max(len(elem) for elem in inputs)
zip(*(islice(cycle(elem), max_length) for elem in inputs))

Demo:

>>> from itertools import islice
>>> inputs = (A, B)  # potentially more
>>> max_length = max(len(elem) for elem in inputs)
>>> zip(*(islice(cycle(elem), max_length) for elem in inputs))
[(70, 1), (60, 2), (50, 3), (40, 4), (30, 1), (20, 2), (10, 3), (0, 4)]

edited Feb 5, 2016 at 2:01

answered Feb 5, 2016 at 1:55

Martijn Pieters

1.1m326 gold badges4.2k silver badges3.4k bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

A.K. Over a year ago

This is perfect. Thank you so much. I tried using the zip() function but I was only receiving pairs up to len(B). I had no idea I could do zip(A, B + B).

Martijn Pieters Over a year ago

@A.K.: zip() terminates at the shortest length. The trick then is to make A the shortest.

Stefan Pochmann Over a year ago

Another way for the arbitrary number case, using islice on the outside of the zip: list(islice(zip(*map(cycle, inputs)), max(map(len, inputs)))).

zondo · Accepted Answer · 2016-02-05 01:56:18Z

10

[(A[x % len(A)], B[x % len(B)]) for x in range(max(len(A), len(B)))]

This will work whether or not A is the larger list. :)

answered Feb 5, 2016 at 1:56

zondo

20.4k8 gold badges50 silver badges89 bronze badges

Comments

John Howard · Accepted Answer · 2016-02-05 01:55:19Z

5

Try using only one for loop instead of two and having the second wrap back to 0 once it gets past its length.

[(A[x], B[x%len(B)]) for x in range(len(A))]

Note that this will only work if A is the longer list. If you know B will always be half the size of A you can also use this:

list(zip(A, B*2))

answered Feb 5, 2016 at 1:55

John Howard

64.9k23 gold badges53 silver badges67 bronze badges

Collectives™ on Stack Overflow

Attempting Python list comprehension with two variable of different ranges

3 Answers 3

3 Comments

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

3 Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related