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.

Solution 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

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

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

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

Solution 5

You can use product from itertools module.

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

View more solutions

Comments

[egor83]

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)

[Li-aung Yip]

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

[MaiaVictor]

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

[Li-aung Yip]

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

[Bi Rico]

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

[Li-aung Yip]

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

[MaiaVictor]

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

[MaiaVictor]

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.

[MaiaVictor]

Just in time for the sprint!

[alan2here]

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