Is there a Python equivalent of range(n) for multidimensional ranges?

Solution 1:

In numpy, it's numpy.ndindex. Also have a look at numpy.ndenumerate.

E.g.

import numpy as np
for x, y in np.ndindex((3,2)):
    print(x, y)

This yields:

0 0
0 1
1 0
1 1
2 0
2 1

Solution 2:

You could use itertools.product():

>>> import itertools
>>> for (i,j,k) in itertools.product(xrange(3),xrange(3),xrange(3)):
...     print i,j,k

The multiple repeated xrange() statements could be expressed like so, if you want to scale this up to a ten-dimensional loop or something similarly ridiculous:

>>> for combination in itertools.product( xrange(3), repeat=10 ):
...     print combination

Which loops over ten variables, varying from (0,0,0,0,0,0,0,0,0,0) to (2,2,2,2,2,2,2,2,2,2).

In general itertools is an insanely awesome module. In the same way regexps are vastly more expressive than "plain" string methods, itertools is a very elegant way of expressing complex loops. You owe it to yourself to read the itertools module documentation. It will make your life more fun.

Solution 3:

There actually is a simple syntax for this. You just need to have two fors:

>>> [(x,y) for x in range(3) for y in range(2)]
[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]

Solution 4:

That is the cartesian product of two lists therefore:

import itertools
for element in itertools.product(range(3),range(2)):
    print element

gives this output:

(0, 0)
(0, 1)
(1, 0)
(1, 1)
(2, 0)
(2, 1)