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.

Answer 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

Answer 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.

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

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

Answer 5

You can use product from itertools module.

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

Answer 6

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 7

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)