Numpy.ndarray：按数组的特定轴进行迭代

Question

我有一个多维数组。 而且我必须通过其某些轴进行迭代。 如果需要所有轴，则可以使用nditer ，但是如果仅需要特定轴，则必须手动执行：

my_array = np.arange(3 * 4 * 5).reshape((3, 4, 5))
for i in range(my_array.shape[0]):
    for j in range(my_array.shape[1]):
        print(i, j)
        # Here should be some processing of the 3rd dimension items of the (i,j)

您不能建议我一个更简单的方法吗？

Answer 1

考虑传递到单个循环，并使用ndindex （ docs ）：

my_array = np.arange(3 * 4 * 5).reshape((3, 4, 5))
for ij in np.ndindex(my_array.shape[:2]):
    i,j=ij
    print(i,j)