Python: Generating all n-length arrays combinations of values within a range

Question

Ok. I'm looking for the smartest and more compact way to do this function

def f():
    [[a,b,c] for a in range(6) for b in range(6) for c in range(6)]

which should generate all the combinations for the values a,b,c like this:

[0,0,0]
[0,0,1]
[0,0,2]
...
[1,0,0]
[1,0,1]
...

and so on...

But I want this to be flexible, so I can change the range or iterable, and also the length of the generated arrays. Range is an easy thing:

def f(min, max):
    [[a,b,c] for a in range(min,max) for b in range(min,max) for c in range(min,max)]

This is ok for 3-length arrays, but I'm thinking now of making 4-length arrays or 7-length arrays and generate all combinations for them in the same range.

It has to exist an easy way, maybe with concatenating arrays or nesting comprehension lists in some way, but my solutions seem to bee too much complex.

Sorry for such a long post.

Answer 1

You can use itertools.product which is just a convenience function for nested iterations. It also has a repeat-argument if you want to repeat the same iterable multiple times:

>>> from itertools import product

>>> amin = 0
>>> amax = 2
>>> list(product(range(amin, amax), repeat=3))
[(0, 0, 0), (0, 0, 1), (0, 1, 0),  (0, 1, 1),  (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]

To get the list of list you could use map:

>>> list(map(list, product(range(amin, amax), repeat=3)))
[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]

However product is an iterator so it's really efficient if you just iterate over it instead of casting it to a list. At least if that's possible in your program. For example:

>>> for prod in product(range(amin, amax), repeat=3):
...     print(prod)  # one example
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(0, 1, 1)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(1, 1, 1)

Answer 2

You can use itertools.product:

from itertools import product

def f(minimum, maximum, n):
    return list(product(*[range(minimum, maximum)] * n))

Drop list to return a generator for memory efficiency.