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

Question

On Python, range(3) will return [0,1,2]. Is there an equivalent for multidimensional ranges?

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


        

Answer 1

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)

Answer 2

You can use product from itertools module.

itertools.product(range(3), range(2))

Answer 3

I would take a look at numpy.meshgrid:

http://docs.scipy.org/doc/numpy-1.6.0/reference/generated/numpy.meshgrid.html

which will give you the X and Y grid values at each position in a mesh/grid. Then you could do something like:

import numpy as np
X,Y = np.meshgrid(xrange(3),xrange(2))
zip(X.ravel(),Y.ravel()) 
#[(0, 0), (1, 0), (2, 0), (0, 1), (1, 1), (2, 1)]

or

zip(X.ravel(order='F'),Y.ravel(order='F')) 
# [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]

Answer 4

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)]

Answer 5

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.

Answer 6

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