简单地遍历 Python 中的多维数组

Question

假设我想遍历多维数组的索引。 我目前拥有的是：

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)

我想减少嵌套并能够在一个简单的单行代码中循环遍历多维数组的所有索引。是不是也可以写成下面的形式？

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

Answer 1

使用 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)

或 如果你不想用这种方式写地图：

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

Answer 2

使用 numpy mgrid，我们可以做：

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)