Is there a Python equivalent of range(n) for multidimensional ranges?
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
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.
There actually is a simple syntax for this. You just need to have two for
s:
>>> [(x,y) for x in range(3) for y in range(2)]
[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]