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

ndindex() is NOT the ND equivalent of range()

(despite some of the other answers here).

It works for your simple example, but it doesn't permit arbitrary start, stop, and step arguments. It only accepts stop, and it hard-codes start to (0,0,...) and hard-codes step to (1,1,...).

Here's an implementation that acts more like the built-in range() function. That is, it permits arbitrary start/stop/step arguments, yet it works on tuples instead of mere integers. Like built-in range(), it returns an iterable.

from itertools import product

def ndrange(start, stop=None, step=None):
    if stop is None:
        stop = start
        start = (0,) * len(stop)
    if step is None:
        step = (1,) * len(start)

    assert len(start) == len(stop) == len(step)    
    for index in product(*map(range, 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)

For numpy users

If your code is numpy-based, then it's probably more convenient to work directly with ndarray objects, rather than an iterable of tuples. It might also be faster. The following implementation is faster than using the above if you planned to convert the result to ndarray.

def ndrange_array(start, stop=None, step=1):
    """
    Like np.ndindex, but accepts start/stop/step instead of
    assuming that start is always (0,0,0) and step is (1,1,1),
    and returns an array instead of an iterator.
    """
    start = np.asarray(start)
    if stop is None:
        stop = start
        start = (0,) * len(stop)

    def ndindex(shape):
        """Like np.ndindex, but returns ndarray"""
        return np.indices(shape).reshape(len(shape), -1).transpose()

    shape = (stop - start + step - 1) // step
    return start + step * ndindex(shape)