CODEWITHSUNDEEP
pythonfor-loopmultidimensional-arrayforeachnumpy-ndarray
Tim Kuipers
Tim Kuipers

Reputation: 1764

Simply loop over multi-dimensional array in Python

Suppose I want to loop through the indices of a multi-dimensional array. What I currently have is:

import numpy as np
points = np.ndarray((1,2,3))
for x in range(points.shape[0]):
    for y in range(points.shape[1]):
        for z in range(points.shape[2]):
            print(x,y,z)

I would like to reduce the nesting and be able to loop over all indices of the multi-dimensional array in a simple one-liner. Could it also be written in the following form?

points = np.ndarray((1,2,3))
for (x, y, z) in ?? :
    print(x,y,z)

Upvotes: 0

Views: 379

Answers (2)

Tim Kuipers
Tim Kuipers

Reputation: 1764

Using numpy mgrid, we can do:

points = np.ndarray((1,2,3))
xs, ys, zs = points.shape
for [x, y, z] in np.mgrid[0:xs, 0:ys, 0:zs].T.reshape(-1,len(points.shape)):
    print(x,y,z)

Upvotes: 0

Pygirl
Pygirl

Reputation: 13349

using iterools:

import itertools
x_dim, y_dim, z_dim = points.shape
for x, y, z in itertools.product(*map(range, (x_dim, y_dim, z_dim))):
    print(x,y,z)

or If you don't wanna use map write in this way:

for x, y, z in itertools.product(range(x_dim), 
                                 range(y_dim), 
                                 range(z_dim)):

0 0 0
0 0 1
0 0 2
0 1 0
0 1 1
0 1 2

Upvotes: 2

Related Questions