numpy迭代数组nditer的实现示例

NumPy 的 nditer 迭代器提供了一种灵活且高效的方式来遍历数组元素。对于初学者来说，它可能看起来有点绕，但掌握了核心用法之后，你会发现它在处理多维数组时比普通 Python 循环顺手得多。下面从最基础的逐个访问开始，一步步拆解。

1 逐个访问数组元素

nditer 最直接的应用就是逐个取出数组中的元素。注意，直接对数组用 for i in a 时，对于多维数组，每次取出的是一行（一维子数组）；而用了 nditer 之后，它会“扁平化”地逐个吐出每个标量元素。

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' 表示列优先（第一维变化最快）。

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

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

---------

0

9

3

12

6

15

1

10

4

13

7

16

2

11

5

14

8

17

---------

0

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

三维时 'F' 的顺序需要稍微留意一下：它先遍历第一维（即 a[0] 和 a[1] 交替），然后才是第二维和第三维。这实际上是 Fortran 语言的内存布局规则。

3 修改数组值

虽说 nditer 可以用来修改数组元素，但通常不推荐这样写——毕竟有更直观的向量化操作。不过了解一下也是有必要的：默认情况下 nditer 把输入数组视为只读，要想修改，必须通过 op_flags=['readwrite'] 告诉它我们要改东西。另外注意，Python 里 x = ... 只是改变了局部变量引用，并不会真的修改数组里的值；必须用 x[...] = ... 这种带省略号索引的方式才能写入。

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 使用外部循环

前面所有例子中，迭代器每次只提供一个元素。这虽然简单，但效率不高——每次循环都要走一遍 Python 解释器，开销不小。更好的做法是把最内层的一维循环交给 NumPy 的原生矢量化操作来处理。nditer 可以通过 flags=['external_loop'] 开启“外部循环”模式，它会把尽可能多的元素打包成一个较大的块（比如整个一维数组）一次性返回。

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]；而列优先时，由于内存布局是 F 顺序，内层循环只能拿到 [0 3]、[1 4]、[2 5] 这样的片段。

[0 1 2 3 4 5]

--------------

[0 3]

[1 4]

[2 5]

5 跟踪索引或多索引

迭代时常常需要知道当前元素的位置索引。比如你想按照内存顺序访问一个数组，但需要同时用 C 顺序、Fortran 顺序或多维索引去查另一个数组。这时 flags 里的 multi_index 和 f_index 就派上用场了。

5.1 multi_index

multi_index 会返回每个元素的多维索引（以元组形式）。下面的例子指定了 order='F'，你可以根据需要调整。

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 < (0, 0)>

3 < (1, 0)>

1 < (0, 1)>

4 < (1, 1)>

2 < (0, 2)>

5 < (1, 2)>

5.2 f_index

如果你只需要一个扁平的线性索引（按 Fortran 顺序），可以用 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>……注意观察规律。

0 < 0>

1 < 2>

2 < 4>

3 < 1>

4 < 3>

5 < 5>

6 缓冲数组元素

开启缓冲模式（buffered）可以让迭代器内部循环一次性处理更大的数据块，大幅减少 Python 解释器的调用次数。下面这个例子强制使用 Fortran 顺序，并同时启用 external_loop 和 buffered，你会发现一次就能拿到所有元素（但顺序仍然是列优先的）。

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

可以看到，a 被广播成了 (2,3)，然后逐元素配对输出。这在很多数值运算中非常实用。