python nested list comprehension for arbitrary number of variables

Question

If I have a tuple of dimensions like dims=(2, 3, 4) and I want to generate a list of all possible indices for this array, a good solution I've found is:

dims = (2, 3, 4)
indices = [range(0, d) for d in dims]
strings = ["[%i][%i][%i]" % (i,j,k) for i in indices[0] for j in indices[1] for k in indices[2]]
print strings:
['[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]']

This is all well and good, but I'm looking for a nice compact way to generate the strings variable where dims is of arbitrary length. I'm trying to avoid a bunch of "if-else" clauses depending on the length of dims. Any ideas?

Answer 1

Use itertools.product:

>>> from itertools import product
>>> dims = (2, 3, 4, 5)
>>> list(product(*map(range, dims)))
[(0, 0, 0, 0), (0, 0, 0, 1), (0, 0, 0, 2), (0, 0, 0, 3), (0, 0, 0, 4), (0, 0, 1, 0), (0, 0, 1, 1), (0, 0, 1, 2), (0, 0, 1, 3), (0, 0, 1, 4), (0, 0, 2, 0), (0, 0, 2, 1), (0, 0, 2, 2), (0, 0, 2, 3), (0, 0, 2, 4), (0, 0, 3, 0), (0, 0, 3, 1), (0, 0, 3, 2), (0, 0, 3, 3), (0, 0, 3, 4), (0, 1, 0, 0), (0, 1, 0, 1), (0, 1, 0, 2), (0, 1, 0, 3), (0, 1, 0, 4), (0, 1, 1, 0), (0, 1, 1, 1), (0, 1, 1, 2), (0, 1, 1, 3), (0, 1, 1, 4), (0, 1, 2, 0), (0, 1, 2, 1), (0, 1, 2, 2), (0, 1, 2, 3), (0, 1, 2, 4), (0, 1, 3, 0), (0, 1, 3, 1), (0, 1, 3, 2), (0, 1, 3, 3), (0, 1, 3, 4), (0, 2, 0, 0), (0, 2, 0, 1), (0, 2, 0, 2), (0, 2, 0, 3), (0, 2, 0, 4), (0, 2, 1, 0), (0, 2, 1, 1), (0, 2, 1, 2), (0, 2, 1, 3), (0, 2, 1, 4), (0, 2, 2, 0), (0, 2, 2, 1), (0, 2, 2, 2), (0, 2, 2, 3), (0, 2, 2, 4), (0, 2, 3, 0), (0, 2, 3, 1), (0, 2, 3, 2), (0, 2, 3, 3), (0, 2, 3, 4), (1, 0, 0, 0), (1, 0, 0, 1), (1, 0, 0, 2), (1, 0, 0, 3), (1, 0, 0, 4), (1, 0, 1, 0), (1, 0, 1, 1), (1, 0, 1, 2), (1, 0, 1, 3), (1, 0, 1, 4), (1, 0, 2, 0), (1, 0, 2, 1), (1, 0, 2, 2), (1, 0, 2, 3), (1, 0, 2, 4), (1, 0, 3, 0), (1, 0, 3, 1), (1, 0, 3, 2), (1, 0, 3, 3), (1, 0, 3, 4), (1, 1, 0, 0), (1, 1, 0, 1), (1, 1, 0, 2), (1, 1, 0, 3), (1, 1, 0, 4), (1, 1, 1, 0), (1, 1, 1, 1), (1, 1, 1, 2), (1, 1, 1, 3), (1, 1, 1, 4), (1, 1, 2, 0), (1, 1, 2, 1), (1, 1, 2, 2), (1, 1, 2, 3), (1, 1, 2, 4), (1, 1, 3, 0), (1, 1, 3, 1), (1, 1, 3, 2), (1, 1, 3, 3), (1, 1, 3, 4), (1, 2, 0, 0), (1, 2, 0, 1), (1, 2, 0, 2), (1, 2, 0, 3), (1, 2, 0, 4), (1, 2, 1, 0), (1, 2, 1, 1), (1, 2, 1, 2), (1, 2, 1, 3), (1, 2, 1, 4), (1, 2, 2, 0), (1, 2, 2, 1), (1, 2, 2, 2), (1, 2, 2, 3), (1, 2, 2, 4), (1, 2, 3, 0), (1, 2, 3, 1), (1, 2, 3, 2), (1, 2, 3, 3), (1, 2, 3, 4)]

Answer 2

Use product from itertools:

In [3]: import itertools

In [4]: dims = (2, 3, 4)

In [5]: for idx in itertools.product(*map(range, dims)):
    ...:     print(idx)
    ...:
(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)

As for dealing with the strings, there are many ways. Here is a quick and dirty way:

In [6]: def block_me(n):
    ...:     return "[" + str(n) + "]"
    ...:

In [7]: for idx in itertools.product(*map(range, dims)):
    ...:     print("".join([block_me(i) for i in idx]))
    ...:
[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]