NumPy 2.x 完全指南【二十九】数组迭代器

墨禹

于 2025-07-29 13:55:40 发布

阅读量618

点赞数 18

CC 4.0 BY-SA版权

分类专栏： NumPy 2.x 完全指南文章标签： numpy python 机器学习

本文链接：https://blog.csdn.net/qq_43437874/article/details/148903658

NumPy 2.x 完全指南专栏收录该内容

30 篇文章

订阅专栏

文章目录

1. 概述
2. flat
3. nditer
4. ndenumerate
5. broadcast

1. 概述

NumPy 数组支持迭代器协议，如果使用 Python for 循环进行迭代时，实际上是在数组的第一个维度（轴 0 ）上进行迭代：

# 一维数组，逐个元素进行
arr = np.array([1, 2, 3])
for x in arr:
    print(x)  # 输出：1 2 3

# 二维数组，将迭代行
arr2d = np.array([[1, 2], [3, 4]])
for row in arr2d:
    print("行:", row)  # 输出：[1,2] 和 [3,4]

所以，NumPy 提供多种高效迭代器，适用于不同场景：

nditer：最核心的多维数组迭代器，支持高效内存访问和广播控制（重点）。
flat：扁平迭代器，将任意维度数组视为一维序列。
ndenumerate：索引迭代器，同时获取元素值及其索引。
broadcast：处理形状不匹配数组的联合迭代。

2. flat

ndarray.flat：数组的一个特殊属性，返回一个扁平迭代器对象（numpy.flatiter），提供了一维视角来访问多维数组的所有元素。

示例 1 ，直接迭代三维数组中的所有元素：

# 创建三维数组
arr = np.arange(8).reshape(2, 2, 2)
print(arr)
# [[[0 1]
#   [2 3]]
#
#  [[4 5]
#   [6 7]]]

# 获取扁平迭代器
flat_iter = arr.flat
print(type(flat_iter))  # <class 'numpy.flatiter'>

# 直接迭代
for element in arr.flat:
    print(element, end=' ')  # 输出: 0 1 2 3 4 5 6 7

示例 2 ，手动控制迭代：

# 手动控制
it = arr.flat
while True:
    try:
        print(next(it), end=' ')  # 0 1 2 3...
    except StopIteration:
        break

3. nditer

NumPy 1.6 版本引入 nditer 迭代器对象，提供了多种灵活的方式，以系统化的方式遍历一个或多个数组的所有元素。

提示：https://numpy.org/doc/stable/reference/generated/numpy.nditer.html#numpy-nditer

3.1 迭代每个元素

使用 nditer 可以完成的最基本任务是访问数组的每个元素。每个元素通过标准的 Python 迭代器接口逐个提供。

示例 1 ，访问数组的每个元素：

a = np.arange(6).reshape(2,3)
for x in np.nditer(a):
    print(x, end=' ') # 0 1 2 3 4 5

3.2 控制迭代顺序

order 参数是 nditer 迭代器中控制元素遍历顺序的核心选项：

C：按行优先顺序访问元素。
F：按列优先顺序访问元素。
A：若所有输入数组均为 Fortran 连续存储，自动用 F 顺序。其他情况（如 C 连续或非连续）默认用 C 顺序。
K（默认）：尽可能贴近元素在内存中的物理存储顺序，不做强制连续性要求。

示例 1 ，默认情况下，为了提高访问效率，元素的访问顺序是根据数组的内存布局来选择的：

arr = np.array([[1, 2],
                [3, 4]], order='F')
for x in np.nditer(arr):
    print(x, end=' ')
# 输出：1 3 2 4

arr = np.array([[1, 2],
                [3, 4]], order='C')
for x in np.nditer(arr):
    print(x, end=' ')
# 输出：1 2 3 4

示例 2 ，通过 order 参数显式指定遍历方向：

arr = np.array([[1, 2], [3, 4]])
# 行优先（C 风格）
for x in np.nditer(arr, order='C'):
    print(x, end=' ')  # 输出：1 2 3 4

# 列优先（Fortran 风格）
for x in np.nditer(arr, order='F'):
    print(x, end=' ')  # 输出：1 3 2 4

3.3 修改元素值

默认情况下，nditer 将输入操作数视为只读对象。若要修改数组元素，必须通过 readwrite 或 writeonly 标识来指定读写模式或只写模式。此时迭代器会创建可修改的缓冲区数组，但修改后的数据不会立即同步到原数组，需在迭代结束后显式触发写回操作。

以下是两种确保数据正确写回的方法：

使用 with 上下文管理器（推荐）：对 nditer，__exit__() 会执行缓冲区的数据回写，并释放资源。
手动调用 close() 方法：显式调用迭代器的 close() 方法会触发缓冲区数据写回原数组，并标记迭代器为不可用状态，调用后再次迭代会报错 StopIteration。

示例 1 ，使用 with 自动管理：

a = np.array([[0, 1], [2, 3]])
with np.nditer(a, op_flags=['readwrite']) as it:
    for x in it:
        x[...] *= 2  # 修改缓冲区数据
# 退出 with 块后自动写回原数组
print(a)  # 输出：[[0, 2], [4, 6]]

示例 2 ，手动调用 close() 方法：

a = np.array([[0, 1], [2, 3]])
it = np.nditer(a, op_flags=['readwrite'])
for x in it:
    x[...] += 1
it.close()  # 手动触发写回
print(a)  # 输出：[[1, 2], [3, 4]]

3.4 外部循环

默认模式下 nditer 迭代器内部处理所有循环逻辑，每次返回单个元素，导致频繁的 Python 层函数调用，效率低下。
通过启用外部循环模式（flags=['external_loop']），将最内层循环移至用户代码中，迭代器每次返回一个连续内存块（一维数组），用户可直接对此块进行操作，减少调用开销。

示例 1 ，对于 C 连续布局的数组来说，元素按行连续存储，因内存完全连续，外部循环时整个数组作为单一数据块返回：

# 元素按行连续存储： [0,1,2,3,4,5]
a = np.arange(6).reshape(2, 3) # [[0, 1, 2], [3, 4, 5]]

for chunk in np.nditer(a, flags=['external_loop'], order='C'):
    print(chunk)  # 输出：[0 1 2 3 4 5]

示例 2 ，对于 F 连续也是内存完全连续，外部循环时整个数组作为单一数据块返回：

# 元素按列连续存储： [1, 3, 2, 4]
a = np.array([[1, 2],
                [3, 4]], order='F')

for chunk in np.nditer(a, flags=['external_loop']):
    print(chunk) # 输出： [1 3 2 4]

示例 3 ，当原数组是 C 连续，在外部循环强制要求按列顺序遍历数组时，会隐式创建数据的临时列优先副本，将每列视为一个连续内存块，每次迭代返回一列的全部元素：、

# 元素按行连续存储： [0,1,2,3,4,5]
a = np.arange(6).reshape(2, 3)
print(a)
# [[0 1 2]
#  [3 4 5]]

for chunk in np.nditer(a, flags=['external_loop'], order='F'):
    print(chunk)
# 输出：
# [0 3]
# [1 4]
# [2 5]

3.5 索引跟踪

flags 参数支持多种索引跟踪标志：

c_index：跟踪 C 顺序（行优先）的扁平索引。
f_index：跟踪 Fortran 顺序（列优先）的扁平索引。
multi_index：跟踪多维索引。

示例 1 ，c_index 和 f_index 是返回内存布局中内存块（一维数组存储）的扁平索引：

a = np.array([[1, 2],
              [3, 4]])  # 内存布局: [1,2,3,4] (行连续)
it = np.nditer(a, flags=['c_index'])

for x in it:
    print(f"[{x}:{it.index}]", end=',')  # [1:0],[2:1],[3:2],[4:3],

示例 2 ，multi_index 则返回逻辑空间上的多维索引：

it = np.nditer(a, flags=['multi_index'])
for x in it:
    print(f"{x}:{it.multi_index}", end=' ')  # 1:(0, 0) 2:(0, 1) 3:(1, 0) 4:(1, 1)

4. ndenumerate

numpy.ndenumerate()：多维索引迭代器，用于同时遍历数组的坐标索引和元素值。

返回值：

返回一个迭代器对象，每次迭代产生一个元组 (index, value)：
- index：当前元素的多维坐标元组 (维度数与输入数组相同)。
- value：当前索引位置的元素值。

示例 1 ，二维数组：

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

for index, value in np.ndenumerate(a):
    print(index, value)
# (0, 0) 1
# (0, 1) 2
# (1, 0) 3
# (1, 1) 4

示例 2 ，三维数组：

# 创建三维数组
b = np.array([[[5,6], [7,8]], [[9,10], [11,12]]])

for pos, val in np.ndenumerate(b):
    print(f"坐标{pos} = {val}")
# 坐标(0, 0, 0) = 5
# 坐标(0, 0, 1) = 6
# 坐标(0, 1, 0) = 7
# 坐标(0, 1, 1) = 8
# 坐标(1, 0, 0) = 9
# 坐标(1, 0, 1) = 10
# 坐标(1, 1, 0) = 11
# 坐标(1, 1, 1) = 12

5. broadcast

numpy.broadcast()：接受两个数组作为输入，创建广播对象并返回一个只读的迭代器，该迭代器返回每个输入序列元素的广播结果的元组。

示例 1 ，迭代广播对象时，遍历每个元素对：

arr1 = np.array([1, 2, 3])   # 形状: (3,)
arr2 = np.array([[4], [5]])  # 形状: (2, 1)

# 创建广播对象
b = np.broadcast(arr1, arr2)

# 广播后形状：(2, 3)
# arr1 扩展为
# [[1, 2, 3],
#  [1, 2, 3]]

# arr2 扩展为：
# [[4, 4, 4],
#  [5, 5, 5]]

for val_x, val_y in b:
    print(val_x + val_y)  # 输出: 5, 6, 7, 6, 7, 8