4.3 迭代为什么慢与nditer

4.3 迭代为什么慢与nditer

本节摘要：直接 for 遍历 ndarray 每一步都发生"装箱"——裸字节被包装成 Python 对象，开销远超计算本身。一维迭代慢十倍起步，多维迭代慢上百倍。正确姿势依次是：向量化、聚合函数、nditer。nditer 以外部循环模式按块迭代，实在无法向量化时它是最后防线。

慢在哪：装箱

第 1 章讲过，Python 变量是带管理信息的对象。for 循环里写 for x in arr，NumPy 每走一步都要在内存里造一个临时 Python 对象把 8 字节裸值装进去，用完再丢。装箱拆箱的成本可以是数值计算本身的几十倍。

import numpy as np
import time

arr = np.arange(1_000_000)

# 姿势一：直接迭代，逐元素装箱
t0 = time.perf_counter()
total = 0.0
for x in arr:
    total += x
t1 = time.perf_counter()
print("for 迭代:", round(t1 - t0, 3), "秒")   # 输出示例：0.107 秒

# 姿势二：内置 sum，同样逐元素
t0 = time.perf_counter()
total = sum(arr)
t1 = time.perf_counter()
print("内置 sum:", round(t1 - t0, 3), "秒")   # 输出示例：0.061 秒

# 姿势三：NumPy 聚合，一次进入C循环
t0 = time.perf_counter()
total = arr.sum()
t1 = time.perf_counter()
print("arr.sum :", round(t1 - t0, 3), "秒")   # 输出示例：0.001 秒
​

百倍差距再次出现。多维更糟：

m = np.arange(1_000_000).reshape(1000, 1000)

# 迭代二维数组，拿到的是一行行的视图，每行再迭代还是装箱
t0 = time.perf_counter()
s = 0
for row in m:
    for x in row:
        s += x
print("双层循环:", round(time.perf_counter() - t0, 3), "秒")
# 输出示例：双层循环: 0.124 秒

t0 = time.perf_counter()
s = m.sum()
print("m.sum    :", round(time.perf_counter() - t0, 3), "秒")
# 输出示例：m.sum    : 0.001 秒
​

处理顺序：向量化优先

遇到"想写循环"的冲动时，按顺序尝试：

import numpy as np

# 想做：把每个元素限制在[0,1]区间
a = np.array([-2.0, 0.3, 1.7])

# 第一选择：已有向量化函数
b = np.clip(a, 0, 1)
print(b)   # [0.  0.3 1. ]

# 第二选择：布尔运算组合（where 三件套）
c = np.where(a < 0, 0, np.where(a > 1, 1, a))
print(c)   # [0.  0.3 1. ]

# 第三选择：apply_along_axis 只是包装循环，救不了性能，别当向量化用
​

nditer：真正的循环加速器

当流程确实无法向量化（比如每步依赖外部状态、调用只能处理标量的库），nditer 比裸 for 快得多。它的关键是 external_loop 模式：每次交给你一整块数据而不是单个元素，装箱次数从百万级降到几百次：

import numpy as np
import time

data = np.arange(1_000_000)

# 裸迭代
t0 = time.perf_counter()
s1 = 0
for x in data:
    s1 += x * 2
t1 = time.perf_counter()
print("裸 for   :", round(t1 - t0, 3), "秒")   # 0.09 秒左右

# nditer 外部循环模式：块级访问
t0 = time.perf_counter()
s2 = 0
for block in np.nditer(data, flags=["external_loop"], op_flags=["readonly"]):
    s2 += int(block.sum())          # 每块用向量化聚合
t1 = time.perf_counter()
print("nditer块 :", round(t1 - t0, 3), "秒")   # 0.002 秒左右
print(s1 == s2)   # True
​

这里聪明之处在于"块内聚合交给 NumPy，块间顺序由你控制"，装箱成本被摊薄到块级别。写数值处理管线时，nditer 还能同时迭代多个数组、声明读写模式，做"就地更新"类流程：

import numpy as np

values = np.array([1.0, -2.0, 3.0, -4.0])

# 就地取绝对值：readwrite 模式直接写回原数组
with np.nditer(values, op_flags=["readwrite"]) as it:
    for v in it:
        v[...] = abs(v)
print(values)   # [1. 2. 3. 4.]
​

注意 v[...] = ... 的写法：迭代变量是零维视图，赋值必须走省略号或 int(v) 取值，直接 v = abs(v) 只会重绑本地名字，什么也改不了——这是 nditer 的头号新手坑。

💡 关键直觉：判断该不该循环的标准很简单——循环体里如果只在做算术，几乎一定能向量化；循环体里如果在做流程控制（状态机、提前退出、IO），才考虑 nditer 或保留循环。

本节要点回顾

• 装箱是原罪：每次迭代元素被包成 Python 对象，成本远超计算本身，一维慢百倍、多维更甚
• 优先级清单：向量化函数、布尔与 where 组合、聚合函数，最后才轮到循环
• nditer 外部循环：块级交付摊薄装箱成本，多数组同步与 readwrite 就地更新是它的杀手锏
• 零维视图坑：nditer 里写回必须用省略号赋值，等号只改本地名
• apply 类伪向量化：只是循环换皮，性能问题原样保留

索引与取数到此告一段落。第 5 章转向"取完怎么摆"：reshape、转置、展平，零拷贝家族的全员亮相。