4.3 数组的迭代

4.3 数组的迭代

NumPy 数组的迭代：深入探索与实践

NumPy 强大的 ndarray 对象不仅支持高效的数值计算，还提供了灵活的迭代方式，允许我们逐个访问和处理数组中的元素。理解和掌握这些迭代技巧对于充分利用 NumPy 的性能至关重要。

4.3.1 基础迭代：使用 for 循环

最简单的迭代方式是使用 Python 的 for 循环。对于一维数组，这非常直接：

import numpy as np
arr = np.array([1, 2, 3, 4, 5])
for element in arr:
    print(element)
​

对于多维数组，直接使用 for 循环会迭代数组的第一个轴（通常是行）。

arr = np.array([[1, 2, 3], [4, 5, 6]])
for row in arr:
    print(row)
​

输出：

[1 2 3]
[4 5 6]
​

这种方式适用于按行处理数据，但如果需要访问单个元素，则需要嵌套循环：

for row in arr:
    for element in row:
        print(element)
​

虽然简单，但这种方法在处理大型数组时效率较低。NumPy 提供了更高效的迭代工具。

4.3.2 nditer：多维数组的通用迭代器

nditer 是 NumPy 提供的一个强大的迭代器对象，用于以各种顺序访问多维数组的元素。它提供了比标准 Python 迭代器更高的灵活性和效率。

基本用法：

arr = np.array([[1, 2, 3], [4, 5, 6]])
for element in np.nditer(arr):
    print(element)
​

nditer 默认以 C 风格（行优先）的顺序迭代数组。

控制迭代顺序：

order 参数可以控制迭代顺序。常用的选项有：

• 'C'：C 风格，行优先。

• 'F'：Fortran 风格，列优先。

arr = np.array([[1, 2, 3], [4, 5, 6]])
for element in np.nditer(arr, order='F'):
    print(element)
​

输出：

1
4
2
5
3
6
​

修改数组元素：

默认情况下，nditer 将数组元素视为只读。要修改数组元素，需要传递 op_flags=['readwrite']。

arr = np.array([[1, 2, 3], [4, 5, 6]])
for element in np.nditer(arr, op_flags=['readwrite']):
    element[...] = element * 2
print(arr)
​

输出：

[[ 2  4  6]
 [ 8 10 12]]
​

使用外部循环：

flags=['external_loop'] 可以将内循环放到 NumPy C 代码中，从而提高速度。它将返回一个一维数组，包含所有元素。

arr = np.array([[1, 2, 3], [4, 5, 6]])
for element in np.nditer(arr, flags=['external_loop'], order='F'):
    print(element)
​

跟踪索引：

使用 flags=['multi_index'] 可以跟踪每个元素的索引。

arr = np.array([[1, 2, 3], [4, 5, 6]])
for element in np.nditer(arr, flags=['multi_index']):
    print(element, arr[element.multi_index])
​

广播迭代：

nditer 还可以用于同时迭代多个数组，实现广播操作。

arr1 = np.array([[1, 2], [3, 4]])
arr2 = np.array([5, 6])
for x, y in np.nditer([arr1, arr2]):
    print(x, y)
​

在这个例子中，arr2 会被广播到与 arr1 相同的形状。

4.3.3 ndenumerate：带索引的迭代

ndenumerate 是另一种迭代器，它返回一个包含索引和元素的元组。

基本用法：

arr = np.array([[1, 2, 3], [4, 5, 6]])
for index, element in np.ndenumerate(arr):
    print(index, element)
​

输出：

(0, 0) 1
(0, 1) 2
(0, 2) 3
(1, 0) 4
(1, 1) 5
(1, 2) 6
​

ndenumerate 比 nditer 更简洁，尤其是在需要索引时。

4.3.4 代码实践对比

下面是一个对比不同迭代方式的示例：

import numpy as np
import time
arr = np.random.rand(1000, 1000)
# 1. 使用 for 循环
start_time = time.time()
for i in range(arr.shape[0]):
    for j in range(arr.shape[1]):
        arr[i, j] *= 2
end_time = time.time()
print(f"For loop time: {end_time - start_time:.4f} seconds")
# 2. 使用 nditer
arr = np.random.rand(1000, 1000)
start_time = time.time()
for element in np.nditer(arr, op_flags=['readwrite']):
    element[...] *= 2
end_time = time.time()
print(f"nditer time: {end_time - start_time:.4f} seconds")
# 3. 使用 ndenumerate
arr = np.random.rand(1000, 1000)
start_time = time.time()
for index, element in np.ndenumerate(arr):
    arr[index] *= 2
end_time = time.time()
print(f"ndenumerate time: {end_time - start_time:.4f} seconds")
# 4. 使用 NumPy 广播 (最快)
arr = np.random.rand(1000, 1000)
start_time = time.time()
arr *= 2
end_time = time.time()
print(f"NumPy broadcasting time: {end_time - start_time:.4f} seconds")
​

通常，NumPy 的广播操作是最快的，因为它利用了底层的优化。nditer 和 ndenumerate 比标准的 Python 循环更快，但仍然不如广播。

4.3.5 迭代流程图

以下是 nditer 的迭代流程图，展示了其主要步骤：

4.3.6 总结

NumPy 提供了多种迭代数组的方式，从简单的 for 循环到更强大的 nditer 和 ndenumerate。选择哪种方式取决于具体的需求和性能考虑。

• 对于简单的迭代，for 循环可能足够。

• 对于需要控制迭代顺序或修改数组元素的情况，nditer 是一个不错的选择。

• 如果需要同时获取索引和元素，ndenumerate 更方便。

• 在可能的情况下，尽量使用 NumPy 的广播操作，以获得最佳性能。

掌握这些迭代技巧可以帮助你更有效地处理 NumPy 数组，并编写出更高效的数值计算代码。