Python: how to get this elegantly?

Question

I'm new to Python and I wrote a function:

def f(x, y, z):
    ret = []
    for i in range(x):
        for j in range(y):
            for k in range(z):
                ret.append((i, j, k))
    return ret


print f(2, 3, 4)

Output:

[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 1, 3), (0, 2, 0), (0, 2, 1), (0, 2, 2), (0, 2, 3), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 0, 3), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 0), (1, 2, 1), (1, 2, 2), (1, 2, 3)]

But I'm not satisfied with that because I think there must be a shorter implementation.

So could any one give me some hint about that?

Well, if your code works and you want to improve it, [Code Reviews SE](https://codereview.stackexchange.com/) would be a better place for your question. Check what's on-topic questions here and there for more details. — Remi Guan, Dec 16 '15 at 12:36

Cory Kramer · Accepted Answer · 2015-12-16T12:50:47.003

8

You can use itertools.product because that is essentially what you are after, the Cartesian product

>>> from itertools import product
>>> list(product(range(2), range(3), range(4)))
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 1, 3), (0, 2, 0), (0, 2, 1), (0, 2, 2), (0, 2, 3), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 0, 3), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 0), (1, 2, 1), (1, 2, 2), (1, 2, 3)]

So replacing your existing function you could do

def f(x, y, z):
    return list(product(range(x), range(y), range(z)))

To remove the number of times you have to type out range, you could accept a single list argument then use a generator expression such as

def f(l):
    return list(product(*(range(i) for i in l)))

So then you could call it as

>>> f([2,3,4])
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 1, 3), (0, 2, 0), (0, 2, 1), (0, 2, 2), (0, 2, 3), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 0, 3), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 0), (1, 2, 1), (1, 2, 2), (1, 2, 3)]

edited Dec 16 '15 at 12:50

answered Dec 16 '15 at 12:37

Cory Kramer

114,268
16
167
218

wow, I had never thought there is such a cool function in python! – Sayakiss Dec 16 '15 at 12:38
@CoryKarmer the best answer.. ..Awesome – Rajarshi Das Dec 16 '15 at 12:40
I find `print list(product(range(2), range(2), range(2), range(2)))` is correct, but I want to eliminate the repeat `range(2)`, so I come to `print list(product([range(2)] * 4))` which output is wrong. Is there any way to eliminate repeat `range(2)` here? – Sayakiss Dec 16 '15 at 12:47
itertools is where the cool stuff happens in python :) – polhemic Dec 16 '15 at 12:50
@Sayakiss See my edit, you could modify your function to take a single list argument, will that more flexible for you? – Cory Kramer Dec 16 '15 at 12:51
@Sayakiss `list(product(*map(range, repeat(2, 4))))` – Brian Rodriguez Dec 16 '15 at 12:52
@CoryKramer Maybe it will be useful in the future. In this case which I'm working with now, `f(x,y,z)` is sufficient. – Sayakiss Dec 16 '15 at 12:53

Praveen · Answer 2 · 2015-12-16T12:46:09.020

You can use list comprehension

 >>> [(i, j, k) for i in range(2) for j in range(3) for k in range(4)]
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 1, 3), (0, 2, 0), (0, 2, 1), (0, 2, 2), (0, 2, 3), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 0, 3), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 0), (1, 2, 1), (1, 2, 2), (1, 2, 3)]

You you can write the function like:

def f(x, y, z):
    return [(i, j, k) for i in range(x) for j in range(y) for k in range(z)]

Python: how to get this elegantly?

2 Answers2