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


        

Answer 1

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)