NumPy 向量化：告别 for 循环

NumPy 之所以快，是因为它把逐元素运算从 Python 解释器搬进了编译好的 C 循环：一次调用处理整个数组，而不是让 Python 去管每个元素。理解了这一点，你会发现”去掉 for 循环”既是性能问题，也是表达方式问题。

向量化不是把循环藏起来，而是把循环交还给底层。你负责描述”对整体做什么”，长度、步长与内存访问交给 NumPy。

1. 为什么 for 循环这么慢

CPython 执行 a[i] * b[i] + 1.0 时要做一长串动作：查名字、取出对象、判断类型、分发乘法协议、装箱结果、再写回。每次下标访问都是一个完整的 Python 调用栈，一千万次循环就是一千万次开销。

1.1 解释器开销与逐元素类型检查

更关键的是 Python 的 list 存的是对象指针，而不是数字本身：每次运算都要把 int/float 拆出数值、算完再重新装箱，类型还要在运行时逐个确认。而 ndarray 存的是连续原始字节，运算时一次确定 dtype，整块交给 C 层循环，中途不创建任何 Python 对象。

import numpy as np

n = 1_000_000
a = np.arange(n, dtype=np.float64)
b = np.arange(n, dtype=np.float64)

# 逐元素循环：每次下标都是一次 Python 对象操作
loop = np.empty_like(a)
for i in range(a.size):
    loop[i] = a[i] * b[i] + 1.0

# 向量化：一次调用，底层 C 循环遍历整块内存
vec = a * b + 1.0

# 两者结果完全一致
assert np.allclose(loop, vec)

同样的计算，循环版本往往比向量化版本慢几十到上百倍。差距不来自算法，而来自”谁在执行循环”。

2. ndarray 的内存布局与 dtype

ndarray 由两部分组成：描述形状、dtype、步长（strides）的元数据，以及一块连续的一维内存缓冲区。所谓”二维数组”，在内存里依然是排成一条直线的字节，只是 NumPy 用 strides 告诉它第 i 行、第 j 列该往后跳多少字节。

dtype 决定解释方式：同样 8 个字节，float64 当作一个双精度浮点，int64 当作一个 64 位整数，误用会静默出错
连续性影响性能：转置或切片得到的数组可能”不连续”，此时某些运算需要先拷贝成连续内存
dtype 影响内存占用：float64 换 float32 直接省一半内存，对大规模中间结果非常可观

这也解释了为什么”向量化更快”的前提是数据放进 ndarray：只有连续、定长的字节块，底层才能做循环展开、SIMD 指令这类优化。

3. 广播：让形状自动对齐

广播（broadcasting）让不同形状的数组直接参与运算，规则只有三条，从最后一维往前比较：维度相等可对齐；维度为 1 可拉伸；缺维度补成 1。只要某一维既不等也不为 1，就直接报错。

m = np.arange(12).reshape(3, 4)   # 形状 (3, 4)

# 列偏移 (4,) 广播成 (3, 4)，每列加同一个数
col = np.array([10, 20, 30, 40])
m + col

# 行偏移 (3, 1) 广播成 (3, 4)，每行加同一个数
row = np.array([[100], [200], [300]])
m + row

# 陷阱：形状 (3,) 与 (4,) 无法广播，直接抛 ValueError
# np.zeros(3) + np.zeros(4)

# 用 None 显式升维，把 (3,) 变成 (3, 1)，结果即外积
outer = np.arange(3)[:, None] * np.arange(4)

最常见的坑是把”一维长度”和”二维某一维”搞混：想按行偏移却写了 (4,)，NumPy 会安静地按列算，结果形状对了但语义完全反了。养成写 [:, None] 或 [None, :] 显式指定方向的习惯，比事后 debug 划算得多。

4. 向量化常用套路

绝大多数显式循环都能改写成下面几类操作之一。掌握它们，等于拿到了把 Python 循环”翻译”成向量化的词典。

4.1 布尔掩码筛选

条件判断用布尔数组一次性表达，然后用它去索引。相比循环里逐个 if，掩码既短又快，还能直接组合多个条件。

temps = np.array([12.5, 31.0, 28.4, 41.7, 19.2])

# 布尔掩码：直接取出满足条件的子集
hot = temps[temps > 30]              # array([31.0, 41.7])

# 多条件用 & | 组合，务必给每个条件加括号
warm = temps[(temps > 20) & (temps < 35)]

# np.where：条件为真取 x，否则取 y，相当于向量化三元表达式
label = np.where(temps > 30, 1, 0)

# 漏斗式筛选：先粗筛子集，再对子集做昂贵运算
score = np.sqrt(temps[temps > 20]) * 2.0

4.2 np.where 与 ufunc

分段函数是循环的重灾区：用 np.where 可嵌套表达多段逻辑，分支更多时 np.select 可读性更好。逐元素数学函数优先用 NumPy 的 ufunc（如 np.sqrt、np.exp、np.clip），它们本身就是 C 实现的逐元素运算，比 math.xxx 配合循环快得多，且天然支持数组与广播。

4.3 np.einsum 与滑窗

涉及矩阵乘法、批量运算或张量收缩时，np.einsum 用下标直接描述运算，语义清晰，还避免生成不必要的中间数组；滑窗统计则用 np.lib.stride_tricks.sliding_window_view 构造零拷贝视图，沿窗口轴聚合即可，无需嵌套循环。

A = np.random.rand(64, 200)
B = np.random.rand(200, 32)

# 下标式描述收缩：等价于 A @ B，语义比 dot 更直观
C = np.einsum("ij,jk->ik", A, B)

# 批量矩阵乘法：一次算完整批，无需遍历 batch 维度
batch = np.random.rand(16, 64, 200)
multi = np.einsum("bij,jk->bik", batch, B)

# 对角线提取也可向量化
sq = np.arange(9).reshape(3, 3)
diag = np.einsum("ii->i", sq)          # array([0, 4, 8])

# 滑窗求移动平均：零拷贝视图 + 沿窗口轴聚合
series = np.random.rand(1000)
win = np.lib.stride_tricks.sliding_window_view(series, 5)
moving_avg = win.mean(axis=1)