Enumerating three variables in python list comprehension

Question

I am trying to print all the possible enumerations of a list for three variables. For example if my input is:

x = 1
y = 1
z = 1

I want the output to be like:

[[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 1, 0], [1, 0, 1], [0, 1, 1], [1, 1, 1]]

If any of the x,y,z variables are higher than 1, it would enumerate all the integers from 0 to the variable value. For example, if x=3 then 0, 1, 2, or 3 would be possible in the first slot of the 3-element lists.

Right now I am creating the list comprehension like this:

output = [ [x,y,z] for x,y,z in range(x,y,z)]

I think something is wrong with the range function?

Do you want the output in exactly that order? And if so, how would be the order for more complex input? — tobias_k
– tobias_k, Commented Jan 18, 2016 at 11:30
I'm not that concerned with order, but would like it to be extensible to non-binary variable values, so like if x=2, there would be lists such as [2,0,0] [2,0,1], [2,1,0], [2,1,1] in the output as well. — William Ross
– William Ross, Commented Jan 18, 2016 at 11:32

gtlambert · Accepted Answer · 2016-01-18 15:44:11Z

15

You could use the product() function from itertools as follows:

from itertools import product

answer = list(list(x) for x in product([0, 1], repeat=3))
print(answer)

Output

[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]

edited Jan 18, 2016 at 15:44

answered Jan 18, 2016 at 11:23

gtlambert

12k2 gold badges32 silver badges48 bronze badges

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

3 Comments

William Ross Over a year ago

Actually, would like the variables possibly be not binary as well, so if x=2, y=1, z=1, there could be lists like [2,1,1] in the output as well... is there a way to do your solution using the x, y, and z variables rather than hard coding 1, and 0 there?

gtlambert Over a year ago

Happy to help, but probably best if you put it in a separate question, as it probably requires a slightly different approach!

tobias_k Over a year ago

Why convert afterwards? Just do the same before converting the product generator to a list, e.g. [list(x) for x in product(...)] or list(map(list, product(...))

tobias_k · Accepted Answer · 2016-01-18 11:35:49Z

10

Complementary to the solutions using product, you could also use a triple list comprehension.

>>> x, y, z = 1, 2, 3
>>> [(a, b, c) for a in range(x+1) for b in range(y+1) for c in range(z+1)]
[(0, 0, 0),
 (0, 0, 1),
 (0, 0, 2),
 (some more),
 (1, 2, 2),
 (1, 2, 3)]

The +1 is necessary since range does not include the upper bound. If you want the output to be a list of lists, you can just do [[a, b, c] for ...].

Note, however, that this will obviously only work is you always have three variables (x, y, z), while product would work with an arbitrary number of lists/upper limits.

answered Jan 18, 2016 at 11:35

tobias_k

83.2k13 gold badges130 silver badges186 bronze badges

Comments

Kasravnd · Accepted Answer · 2016-01-18 15:34:13Z

6

You can use range() function within a list comprehension and itertools.product function:

>>> x = 1
>>> y = 1
>>> z = 1
>>> from itertools import product
>>> list(product(*[range(i+1) for i in [x,y,z]]))
[(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]

This approach will work for different numbers too:

>>> x = 2
>>> y = 2
>>> z = 2
>>> 
>>> list(product(*[range(i+1) for i in [x,y,z]]))
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 2, 0), (0, 2, 1), (0, 2, 2), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 2, 0), (1, 2, 1), (1, 2, 2), (2, 0, 0), (2, 0, 1), (2, 0, 2), (2, 1, 0), (2, 1, 1), (2, 1, 2), (2, 2, 0), (2, 2, 1), (2, 2, 2)]
>>>

edited Jan 18, 2016 at 15:34

answered Jan 18, 2016 at 11:28

Kasravnd

108k19 gold badges167 silver badges196 bronze badges

2 Comments

William Ross Over a year ago

Is there a way to get the output in the form of lists rather than tuples?

Kasravnd Over a year ago

@WilliamRoss Yep you can use a list comprehension instead of list function. [list(item) for item in product(*[range(i+1) for i in [x,y,z]])]

Ali · Accepted Answer · 2016-01-18 12:30:25Z

4

If you need it in form of a list of lists (instead of list of tuples) you can use map over the output of the answer by Kasramvd, that is:

map(list,list(product(*[range(i+1) for i in [x,y,z]])))

answered Jan 18, 2016 at 12:30

Ali

1,6351 gold badge13 silver badges19 bronze badges

1 Comment

tobias_k Over a year ago

1) in Python 3, map produces a generator, not a list, and 2) no need to turn the generator returned by product into a list before passing it to map. So this should rather be list(map(list, product(...))

Collectives™ on Stack Overflow

Enumerating three variables in python list comprehension

4 Answers 4

3 Comments

Comments

2 Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

3 Comments

Comments

2 Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related