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.

Felipa

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

Ioannina

+1: The syntax for that is alarmingly similar to what the OP originally asked for. Well played! – Scanlan

As Li-aung pointed this is alarmingly similar to what I asked for, so it is, undoubtedly, the best answer to the topic. – Felipa

Li-aung Yip answer is great, too, and has some learning on it as it shows the cartesian product can be used for the same purpose. – Felipa

This makes loads of things with Numpy Arrays gigantically more generic, soo powerful. – Yoong

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.

Scanlan

just a tiny improvement over your last answer: for c in product(*([xrange(5)]*3)): print c: from (0,0,0) to (4,4,4) – Malaysia

It's actually better to use itertools.tee() if you want exact replicas - I believe the underlying implementation is more efficient due to caching. – Scanlan

@agf: good catch - evidently time for me to sleep. Edited to that effect. – Scanlan

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

Albertype

This is good, but I would like to point that it can get a little verbose: for (x,y) in [(x,y) for x in range(3) for y in range(2)]: – Felipa

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)

Ramirez

Answer 5

You can use product from itertools module.

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

Reproachless

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

Lamprophyre

it would be good to also mention numpy.mgrid and numpy.ogrid here. – Hesterhesther

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)

Idellaidelle

Just in time for the sprint! – Felipa

good job for providing a full answer. when most of the answers are focusing only on the single use case that OP mentioned, you actually provided an answer for the question of equivalent range function for multidimensional case, which covers all the use cases – Lautrec

Answer 8

arr = np.array(range(1,7))

print(np.reshape(arr,(2,3)))

[[1 2 3]
 [4 5 6]]

Krishnakrishnah

Thank you for your interest in contributing to the Stack Overflow community. This question already has quite a few answers—including one that has been extensively validated by the community. Are you certain your approach hasn’t been given previously? If so, it would be useful to explain how your approach is different, under what circumstances your approach might be preferred, and/or why you think the previous answers aren’t sufficient. Can you kindly edit your answer to offer an explanation? – Tonguing