一维数组的遍历

一维数组的很多操作都和Python中的序列类型相同，其迭代可以直接使用for来完成遍历。

  1. import numpy as np
  3. a = np.arange(5)
  4. for i in a:
  5.     print(i, end=' ')   # 0 1 2 3 4

多维数组的遍历

对于多维数组，不能简单地使用for来进行遍历，通过ndarray.flat，可以将多维数组「降维」，从而简单完成遍历。

  1. import numpy as np
  3. a = np.array([[1, 2, 3], [6, 7, 8]])
  5. for i in a:
  6.     print(i)
  8. # [1 2 3]
  9. # [6 7 8]
  11. """
  12. 使用 for 进行遍历，它只处理了第一个维度，并不是按元素进行的遍历操作
  13. 如果我们知道数组的维度情况，就可以采用多个 for 循环嵌套来完成元素遍历
  14. 或者采用递归的写法
  15. """
  16. def my_traverse(array):
  17.     for i in array:
  18.         if isinstance(i, np.ndarray):
  19.             my_traverse(i)
  20.         else:
  21.             print(i)

ndarray.flat返回了数组对象的迭代器，通过这个迭代器，就可以逐元素遍历。在 「Learning Python」中，已经详细讲解了迭代器和迭代协议，下面直接通过例子说明

  1. a = np.array([[1, 2, 3], [6, 7, 8]])
  3. iter = a.flat
  4. print(type(iter))   # <class 'numpy.flatiter'>
  6. print(next(iter))   # 1
  7. print(next(iter))   # 2
  9. for data in iter:
  10.     print(data, end=' ')    # 3 6 7 8

flat有「平整」的意思，这里可以理解为，将数组摊平到了一个维度上。

  1. a = np.array([[1, 2, 3], [6, 7, 8]])
  3. """
  4. 就像操作一维数组一样
  5. """
  6. print(a.flat[3], a.flat[4]) # 6 7
  7. a.flat[0] = 666
  9. print(a)
  10. # [[666   2   3]
  11. #  [  6   7   8]]

更通用的遍历方法

numpy.nditer是一种更加灵活的迭代器对象。

  1. import numpy as np
  3. a = np.arange(6).reshape(2, 3)
  5. print(a)
  6. # [[0 1 2]
  7. #  [3 4 5]]
  9. for i in np.nditer(a):
  10.     print(i, end=' ')   # 0 1 2 3 4 5

控制遍历顺序

我们可以选择遍历数组对象的顺序方式。默认的情况是根据数组元素在内存中的存储「模式」而定的，这样做更能提升访问效率，而并不是我们直观上的先遍历行元素，再遍历列元素。

  1. b = a.T     # 转置
  2. print(b)
  3. # [[0 3]
  4. #  [1 4]
  5. #  [2 5]]
  7. """
  8. 遍历结果和 a 的是一样的，说明通过转置计算得到的 b
  9. 在内存上的元素存储「模式」是一样的
  10. """
  11. for i in np.nditer(b):
  12.     print(i, end=' ')   # 0 1 2 3 4 5

比较常用的遍历顺序有C-order和F-order，前者是我们常见的行序优先，后者则是列序优先，现在用这两种方式来遍历b数组，再和flat方式做个对比：

  1. for i in np.nditer(b, order='C'):   # C style
  2.     print(i, end=' ')   # 0 3 1 4 2 5
  4. for i in np.nditer(b, order='F'):   # Fortran style
  5.     print(i, end=' ')   # 0 1 2 3 4 5
  7. for i in b.flat:
  8.     print(i, end=' ')   # 0 3 1 4 2 5 说明 flat 是按照 C style 将元素「摊平」的

修改元素值

除了指定遍历元素的方式以外，还可以在遍历的过程中对数组元素进行修改。默认的遍历操作都是只读的，可以通过readwrite和writeonly将操作设定为读写或者只读模式。

nditer首先会生成可写的数组缓冲区，然后将修改后的元素值拷贝回原来的数组。因此，我们必须要告知迭代结束，以便后续元素拷贝。这里使用上下文管理器：

  1. with np.nditer(a, op_flags=['readwrite']) as it:
  2.     for i in it:
  3.         i[...] = i + 1  # 修改元素值
  5. print(a)
  6. # [[1 2 3]
  7. #  [4 5 6]]

i[...]的写法在「ndarray对象」那一小结提到过，目的就是为了修改它的值。这里感觉上可能会有些矛盾，因为给变量的赋值，只是将变量指向了新的地方，并未在原来内存空间上进行修改，所以直接常规的赋值方式并不能更改数组中的元素值。

在上面的例子中，因为i本质上是一个零维数组，通过上面的方法才能够正确修改其值。参考：What does “three dots” in Python mean when indexing what looks like a number?

看看使用flat迭代修改会怎么样：

  1. import numpy as np
  3. a = np.arange(6).reshape(2, 3)
  5. for i in a.flat:
  6.     i += 3      # 这里 i 的本质是整数，无法在原处修改值
  7. print(a)
  9. # [[0 1 2]
  10. #  [3 4 5]]

易错点

在使用 Numpy 数组进行遍历更新时，可能凭直觉，我们都会写出这样的代码：

  1. import numpy as np
  3. a = np.arange(6).reshape(2, 3)
  5. for i in a:
  6.     i += 3        # 为什么不写 i = i + 3 我就不解释了
  8. print(a)
  10. # [[3 4 5]
  11. #  [6 7 8]]

哟，完蛋，一不小心对数组进行了遍历， 而且把每个元素的值更新成功了。既然如此，为什么前面讲了一堆很复杂的东西呢？

虽然这里结果是正确的，但是程序并不是按照我们所想的逻辑运行出来的——它并没有对每一个元素进行遍历。在最开头，我们就知道了，直接使用 for 来遍历，程序只会对数组的第一维进行处理。

仔细思考下，循环里的 i 表示的是什么？

可以插入代码试验一下：

  1. a = np.arange(6).reshape(2, 3)
  3. for i in a:
  4.     print(i)
  6. # <class 'numpy.ndarray'>
  7. # <class 'numpy.ndarray'>

首先，我们知道，数组只被访问了两次；其次，我们操作的对象 i ，它是 Numpy 子数组。所以，上面加法的本质是这样的：

  1. # 第一次访问
  2. i = [0, 1, 2]
  3. i += 3
  4. # 第二次访问
  5. i = [3, 4, 5]
  6. i += 3

那为啥结果是对的？这就是前面讨论过的 Numpy 的广播机制 。 3 被「扩展」成了 [3, 3, 3] 来和子数组相加。所以， Numpy 数组的遍历操作没有我们想的那么简单，如果我们的需求仅仅是每个元素 +3 这么简单，当然可以这么来处理，但是要知道其背后的原理。

如果要处理更复杂的元素遍历修改操作（比如，奇数元素+1；偶数元素-1），就还是要借助迭代器这种比较复杂的方法来处理。