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

PythonNumpyIterationRange

Python Problem Overview


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.

Python Solutions


Solution 1 - Python

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

Solution 2 - Python

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.

Solution 3 - Python

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

Solution 4 - Python

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)

Solution 5 - Python

You can use product from itertools module.

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

Solution 6 - Python

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

Solution 7 - Python

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)

Attributions

All content for this solution is sourced from the original question on Stackoverflow.

The content on this page is licensed under the Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) license.

Content TypeOriginal AuthorOriginal Content on Stackoverflow
QuestionMaiaVictorView Question on Stackoverflow
Solution 1 - PythonJoe KingtonView Answer on Stackoverflow
Solution 2 - PythonLi-aung YipView Answer on Stackoverflow
Solution 3 - PythondhgView Answer on Stackoverflow
Solution 4 - Python0x90View Answer on Stackoverflow
Solution 5 - PythonPraveen GollakotaView Answer on Stackoverflow
Solution 6 - PythonJoshAdelView Answer on Stackoverflow
Solution 7 - PythonStuart BergView Answer on Stackoverflow