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)]
So, for example, looping though the tiles of a rectangular area on a tile-based game could be written as:
for x,y in range((3,2)):
Note I’m not asking for an implementation. I would like to know if this is a recognized pattern and if there is a built-in function on Python or it’s standard/common libraries.
Answers:
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)]
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.
You can use product from itertools module.
itertools.product(range(3), range(2))
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)
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)]
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
Numpy’s ndindex() works for the example you gave, but it doesn’t serve all use cases. Unlike Python’s built-in range(), which permits both an arbitrary start, stop, and step, numpy’s np.ndindex() only accepts a stop. (The start is presumed to be (0,0,...), and the step is (1,1,...).)
Here’s an implementation that acts more like the built-in range() function. That is, it permits arbitrary start/stop/step arguments, but it works on tuples instead of mere integers.
import sys
from itertools import product, starmap
# Python 2/3 compatibility
if sys.version_info.major < 3:
from itertools import izip
else:
izip = zip
xrange = range
def ndrange(start, stop=None, step=None):
if stop is None:
stop = start
start = (0,)*len(stop)
if step is None:
step = (1,)*len(stop)
assert len(start) == len(stop) == len(step)
for index in product(*starmap(xrange, izip(start, stop, step))):
yield index
Example:
In [7]: for index in ndrange((1,2,3), (10,20,30), step=(5,10,15)):
...: print(index)
...:
(1, 2, 3)
(1, 2, 18)
(1, 12, 3)
(1, 12, 18)
(6, 2, 3)
(6, 2, 18)
(6, 12, 3)
(6, 12, 18)
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)]
So, for example, looping though the tiles of a rectangular area on a tile-based game could be written as:
for x,y in range((3,2)):
Note I’m not asking for an implementation. I would like to know if this is a recognized pattern and if there is a built-in function on Python or it’s standard/common libraries.
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)]
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.
You can use product from itertools module.
itertools.product(range(3), range(2))
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)
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)]
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
Numpy’s ndindex() works for the example you gave, but it doesn’t serve all use cases. Unlike Python’s built-in range(), which permits both an arbitrary start, stop, and step, numpy’s np.ndindex() only accepts a stop. (The start is presumed to be (0,0,...), and the step is (1,1,...).)
Here’s an implementation that acts more like the built-in range() function. That is, it permits arbitrary start/stop/step arguments, but it works on tuples instead of mere integers.
import sys
from itertools import product, starmap
# Python 2/3 compatibility
if sys.version_info.major < 3:
from itertools import izip
else:
izip = zip
xrange = range
def ndrange(start, stop=None, step=None):
if stop is None:
stop = start
start = (0,)*len(stop)
if step is None:
step = (1,)*len(stop)
assert len(start) == len(stop) == len(step)
for index in product(*starmap(xrange, izip(start, stop, step))):
yield index
Example:
In [7]: for index in ndrange((1,2,3), (10,20,30), step=(5,10,15)):
...: print(index)
...:
(1, 2, 3)
(1, 2, 18)
(1, 12, 3)
(1, 12, 18)
(6, 2, 3)
(6, 2, 18)
(6, 12, 3)
(6, 12, 18)