NumPy 迭代数组的几种方法详解

NumPy 提供了多种迭代数组的方式，适用于不同场景（如逐元素操作、按行/列遍历、高性能计算等）。以下是常见的迭代方法及其适用场景：

---

一、基础迭代方法

1. 一维数组：直接迭代元素

一维数组的迭代与 Python 列表类似，直接遍历每个元素：

import numpy as np

arr = np.array([1, 2, 3, 4])
for element in arr:
    print(element)  # 依次输出：1, 2, 3, 4

2. 多维数组：默认迭代最外层轴（行）

多维数组默认按最外层轴（行）​迭代，返回的是一维子数组（行向量）：

arr_2d = np.array([[1, 2], [3, 4], [5, 6]])  # 3行2列

# 默认迭代行（最外层轴）
for row in arr_2d:
    print(row)
# 输出：
# [1 2]  # 第1行
# [3 4]  # 第2行
# [5 6]  # 第3行

---

二、逐元素迭代

1. np.nditer()（推荐，高效通用）

np.nditer 是 NumPy 提供的多维迭代器，可高效遍历任意维度的数组元素，支持控制遍历顺序和数据类型。

基本用法（逐元素访问）

arr = np.array([[1, 2], [3, 4]])

# 默认按行优先（C风格）顺序迭代所有元素
for element in np.nditer(arr):
    print(element)  # 依次输出：1, 2, 3, 4

控制遍历顺序

• ​**order='C'**​（默认，行优先，C风格）：按行依次遍历（内存连续，效率高）。
• ​**order='F'**​（列优先，Fortran风格）：按列依次遍历。

arr = np.array([[1, 2], [3, 4]])

# 列优先遍历（F风格）
for element in np.nditer(arr, order='F'):
    print(element)  # 依次输出：1, 3, 2, 4

同时获取元素索引（通过 flags 和 multi_index）

若需知道当前元素的索引位置，启用 flags=['multi_index'] 并通过 it.multi_index 获取：

arr = np.array([[1, 2], [3, 4]])

it = np.nditer(arr, flags=['multi_index'])
while not it.finished:
    idx = it.multi_index  # 当前元素的索引（如 (0,0)、(0,1)）
    print(f"索引 {idx} 的值: {arr[idx]}")
    it.iternext()  # 移动到下一个元素
# 输出：
# 索引 (0, 0) 的值: 1
# 索引 (0, 1) 的值: 2
# 索引 (1, 0) 的值: 3
# 索引 (1, 1) 的值: 4

原地修改元素（通过 op_flags）

若需在迭代中修改原数组，需设置 op_flags=['readwrite'] 并通过 element[...] 赋值：

arr = np.array([1, 2, 3])
for element in np.nditer(arr, op_flags=['readwrite']):
    element[...] = element * 2  # 每个元素乘以2
print(arr)  # 输出：[2 4 6]

---

2. np.ndenumerate()（获取索引和值）

若需同时获取元素的索引和值，可使用 np.ndenumerate，返回 (索引元组, 值) 的迭代器：

arr = np.array([[1, 2], [3, 4]])

for idx, value in np.ndenumerate(arr):
    print(f"索引 {idx} 的值: {value}")
# 输出：
# 索引 (0, 0) 的值: 1
# 索引 (0, 1) 的值: 2
# 索引 (1, 0) 的值: 3
# 索引 (1, 1) 的值: 4

---

三、按轴迭代（行/列/维度）

1. 遍历特定轴（如按行或列）

通过 axis 参数控制遍历的维度，结合数组切片实现：

按行遍历（默认）

arr_2d = np.array([[1, 2], [3, 4], [5, 6]])

for row in arr_2d:  # 默认按行（axis=0）
    print(row)  # 每行是一个一维数组

按列遍历（axis=1）

需通过转置或索引切片实现：

# 方法1：转置后按行遍历（等价于原数组的列）
for col in arr_2d.T:  # 转置后，原列变为行
    print(col)  # 依次输出：[1 3 5], [2 4 6]

# 方法2：直接索引列（需嵌套循环）
for j in range(arr_2d.shape[1]):  # 遍历列数
    column = arr_2d[:, j]  # 所有行的第j列
    print(column)  # 依次输出：[1 3 5], [2 4 6]

2. 遍历更高维度的轴

对于三维及以上数组，通过指定轴的索引范围迭代：

arr_3d = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])  # 形状 (2,2,2)

# 遍历第一个轴（如2个2x2矩阵）
for matrix in arr_3d:
    print(matrix)  # 每个matrix是2x2的二维数组

---

四、性能对比与推荐

方法	适用场景	特点
​直接迭代​	一维数组或仅需最外层元素（如行）	简单直接，但多维数组默认仅迭代最外层轴（行）。
​**np.nditer**​	高效逐元素遍历（推荐）	支持多维、控制遍历顺序（C/F）、原地修改、获取索引，性能最优。
​**np.ndenumerate**​	需同时获取索引和值	比手动索引更简洁，适合需要索引的场景。
​按轴切片迭代​	按特定维度（如列）遍历	需结合转置或索引操作，灵活性高但代码稍复杂。

​推荐优先级​：

• ​逐元素操作​ → 用 np.nditer（尤其需高性能或修改原数组时）；
• ​需索引和值​ → 用 np.ndenumerate；
• ​按行/列遍历​ → 直接迭代或转置后处理。

---

五、注意事项

1. ​避免显式循环​：NumPy 的向量化操作（如 arr * 2）比迭代更快，仅在必须逐元素处理时使用迭代；
2. ​**nditer 的 readwrite 模式**​：修改元素时需用 element[...] 而非 element = ...（后者不改变原数组）；
3. ​内存连续性​：行优先（C风格）遍历通常比列优先（F风格）更快（因内存布局连续）。

通过合理选择迭代方法，可以兼顾代码简洁性、可读性与性能！ 🚀