如何使用zip遍历ndarray全部元素？解决遍历引发的IndexError

嘿，咱们先弄明白你为啥会碰到这个IndexError：
当你调用zip(a, b)的时候，a和b都是二维ndarray对象，zip会把它们的每一行作为一组配对起来。举个例子，第一次循环时，i和j分别是a[0]和b[0]——也就是一维数组array([1,2,3,4])。这时候你尝试写i[0][0]，就相当于先取i[0]得到标量1，再对这个标量去索引[0]，这肯定会报错啊！

接下来给你几个能实现需求的方案：

方案1：先展平数组再遍历 #

如果你只是想逐个获取所有元素，最直接的方式是把二维数组展平成一维迭代器，再用zip配对（虽然这里b是a的引用，元素完全一样，但这个方案也适用于a和b是不同数组的场景）：

import numpy as np

a = np.array([[1,2,3,4],[5,6,7,8]])
b = a

# 利用flat属性获取数组的扁平迭代器，逐个取元素
for elem_a, elem_b in zip(a.flat, b.flat):
    # 因为elem_a和elem_b值相同，这里输出elem_a即可
    print(elem_a)

要是你需要输出类似a[0][0] = 1这种带索引的格式，可以结合行列循环来实现：

rows, cols = a.shape
for i in range(rows):
    for j in range(cols):
        print(f"a[{i}][{j}] = {a[i][j]}")

方案2：用numpy的ndenumerate遍历索引与元素 #

如果你想同时拿到元素的坐标和值，numpy.ndenumerate是个不错的选择，它会返回每个元素的索引元组和对应的值：

for (i,j), val in np.ndenumerate(a):
    print(f"a[{i}][{j}] = {val}")

方案3：嵌套zip遍历（适合需要同时处理行和元素的场景） #

如果你坚持要用zip来实现，那可以嵌套两层zip：外层配对行，内层配对每行里的元素：

for row_a, row_b in zip(a, b):
    # 内层zip把每行的元素一一配对
    for elem_a, elem_b in zip(row_a, row_b):
        print(elem_a)

这样就能完美实现逐个输出从a[0][0]到a[1][3]所有元素的需求啦！

内容的提问来源于stack exchange，提问作者Srivatsan