NumPy 数组迭代器 nditer 使用详解：从基础到进阶

1. 逐个访问数组元素

使用 nditer 可以逐个访问数组中的元素。对于一维数组，直接使用 for 循环与 nditer 的结果相同，但对于多维数组，两者存在明显区别。

import numpy as np
a = np.arange(4).reshape(2, 2)
for i in a:
    print(i)
for i in np.nditer(a):
    print(i)

输出结果：

[0 1]
[2 3]
0
1
2
3

可以看到，直接遍历二维数组得到的是每一行的一维数组，而 nditer 则逐个输出所有元素。

2. 控制迭代顺序

nditer 对象提供了 order 参数来控制迭代顺序，可选值包括 'K'、'C' 和 'F'。默认值为 'K'，表示保持数组原有的内存顺序；'C' 表示 C 顺序（行优先）；'F' 表示 Fortran 顺序（列优先）。

2.1 二维数组示例

a = np.arange(6).reshape(2, 3)
print(a)
for i in np.nditer(a, order='C'):
    print(i)
print("---------")
for i in np.nditer(a, order='F'):
    print(i)
print("---------")
for i in np.nditer(a, order='K'):
    print(i)

输出结果：

[[0 1 2]
 [3 4 5]]
0
1
2
3
4
5
---------
0
3
1
4
2
5
---------
0
1
2
3
4
5

在二维情况下，'F' 参数确实可以看作按列读取元素。

2.2 三维数组示例

a = np.arange(18).reshape(2, 3, 3)
print(a)
for i in np.nditer(a, order='C'):
    print(i)
print("---------")
for i in np.nditer(a, order='F'):
    print(i)
print("---------")
for i in np.nditer(a, order='K'):
    print(i)

输出结果（部分）：

[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]]
0
1
2
...
17
---------
0
9
3
12
6
15
1
10
4
13
7
16
2
11
5
14
8
17
---------
0
1
2
...
17

注意三维数组在列优先（'F'）时的读取顺序。

3. 修改数组值

默认情况下，nditer 将输入数组视为只读对象。要修改数组元素，必须指定读写或只写模式，这通过 op_flags 参数控制。Python 中的常规赋值只会更改局部或全局变量字典中的引用，并不会修改现有变量。因此，要实际修改数组的元素，应使用省略号索引（...）。

a = np.arange(6).reshape(2, 3)
for x in np.nditer(a, op_flags=['readwrite']):
    x[...] = 2 * x
print(a)

输出结果：

[[ 0  2  4]
 [ 6  8 10]]

4. 使用外部循环

在之前的示例中，元素由迭代器一次提供一个，虽然简单方便，但效率不高。更好的方法是将一维最内层循环移动到迭代器外部的代码中，这样 NumPy 的矢量化操作可以作用于被访问元素的较大块。

nditer 会尝试提供尽可能大的内部循环块。通过强制 'C' 和 'F' 顺序，可以得到不同的外部循环大小。通过指定迭代器标志来启用此模式。

a = np.arange(6).reshape(2, 3)
for x in np.nditer(a, flags=['external_loop'], order='C'):
    print(x)
print("--------------")
for x in np.nditer(a, flags=['external_loop'], order='F'):
    print(x)

输出结果：

[0 1 2 3 4 5]
--------------
[0 3]
[1 4]
[2 5]

注意行序优先时输出的是一个数组，而不是一个个的数字元素。

5. 跟踪索引或多索引

在迭代期间，你可能希望在计算中使用当前元素的索引。例如，你可能希望按内存顺序访问数组的元素，然后使用 C 顺序、Fortran 顺序或多维索引来查找不同数组中的值。

5.1 multi_index

下面的代码中，order 参数可根据需要选择，也可以不写（默认按行逐个读取元素）。

a = np.arange(6).reshape(2, 3)
it = np.nditer(a, flags=['multi_index'], order="F")
while not it.finished:
    print("%d < %s>" % (it[0], it.multi_index))
    it.iternext()

输出结果：it.multi_index 代表元素的索引，以元组形式输出。

0 < (0, 0)>
3 < (1, 0)>
1 < (0, 1)>
4 < (1, 1)>
2 < (0, 2)>
5 < (1, 2)>

5.2 f_index

a = np.arange(6).reshape(2, 3)
it = np.nditer(a, flags=['f_index'])
while not it.finished:
    print("%d < %d>" % (it[0], it.index))
    it.iternext()

输出结果：索引的编号，以列序优先。

0 < 0>
1 < 2>
2 < 4>
3 < 1>
4 < 3>
5 < 5>

6. 缓冲数组元素

通过启用缓冲模式，迭代器提供给内部循环的块可以变得更大，从而显著减少 Python 解释器的开销。在强制 Fortran 迭代顺序的示例中，当启用缓冲时，内部循环可以一次性查看所有元素。

a = np.arange(6).reshape(2, 3)
for x in np.nditer(a, flags=['external_loop', 'buffered'], order='F'):
    print(x)

输出结果：

[0 3 1 4 2 5]

7. 广播数组迭代

当迭代多个数组时，nditer 会自动处理广播规则，将较小的数组广播到较大的形状。

a = np.arange(3)
b = np.arange(6).reshape(2, 3)
for x, y in np.nditer([a, b]):
    print("%d:%d" % (x, y))

输出结果：

0:0
1:1
2:2
0:3
1:4
2:5