4. 向量化基础

在 NumPy 中，向量化并不是一个抽象口号，而是由两类具体机制支撑起来的：通用函数（ufunc）负责把标量层面的运算提升为对整个数组的逐元素运算，广播（broadcasting）负责让形状不同的数组在规则允许的范围内自动对齐。理解这两者，才能把原本写成 Python 循环的代码安全、高效地改写为数组表达式。

一、ufunc：逐元素运算的底层引擎

通用函数是 NumPy 中对 ndarray 执行逐元素操作的函数对象，例如 np.add、np.multiply、np.sqrt、np.exp、np.maximum、np.sin 等。它们有两个关键特征：一是逐元素，输入数组的每个元素独立参与计算，输出与输入形状一致（除非使用 out 或归约参数）；二是类型化，同一个 ufunc 会根据输入 dtype 选择对应的底层循环（如 float64 与 int32 走不同实现），并在必要时按 NumPy 的类型提升规则确定输出 dtype。

日常写 a + b、a * 2、np.sqrt(a) 时，背后调用的正是 ufunc。运算符重载只是语法糖，真正的计算发生在 C 层的循环里，因此避免了 Python 解释器对每个元素逐一调度，这正是向量化提速的主要来源。ufunc 还支持几个实用参数：out= 可把结果写入预先分配的数组，减少临时对象；where= 可配合掩码只对部分位置计算；dtype= 可显式指定计算类型，避免意外的类型提升。

import numpy as np

a = np.array([1.0, 4.0, 9.0, 16.0])
b = np.array([2.0, 2.0, 3.0, 4.0])

# 逐元素开方与相加，均为 ufunc 调用
r1 = np.sqrt(a)              # [1. 2. 3. 4.]
r2 = np.add(a, b)            # [ 3.  6. 12. 20.]
r3 = np.maximum(a, b)        # [ 2.  4.  9. 16.]

# 用 out 复用缓冲区，避免反复分配
buf = np.empty_like(a)
np.multiply(a, b, out=buf)   # buf = [ 2.  8. 27. 64.]

二、广播：形状不同的数组如何对齐

广播解决的是“形状不一致但语义上可以逐元素运算”的问题。规则可以概括为三步：先把两个数组的形状右对齐；然后从最右侧维度开始逐维比较；每一维上，若两者相等，或其中一方为 1，则该维兼容，结果为该维的较大值；若既不相等也没有 1，则抛出 ValueError。形状为 1 的维度会被“拉伸”成对方长度，但不复制数据，只是逻辑上重复参与计算，因此内存开销很小。

例如形状 (3, 4) 与 (4,) 相加：右对齐后第二维都是 4，第一维一方缺失视为 1，于是 (4,) 被广播到每一行，得到 (3, 4)。再如 (3, 1) 与 (1, 4) 相加，两个长度为 1 的维度分别被拉伸，结果为 (3, 4)。而 (3, 4) 与 (3,) 相加则会失败，因为右对齐后第一维是 4 与 3，互不相等且都不为 1。

三、把显式循环改写成向量表达式

改写的基本思路是：先识别循环体在“每个元素上做了什么”，再判断这些操作能否用 ufunc 表达，最后检查参与运算的数组形状是否满足广播规则。下面用一个常见的场景说明——对一批样本做标准化，即每个特征减去均值再除以标准差。

# 循环写法：慢且冗长
X = np.random.rand(1000, 5)
mu = X.mean(axis=0)      # 形状 (5,)
sd = X.std(axis=0)       # 形状 (5,)
Z = np.empty_like(X)
for i in range(X.shape[0]):
    for j in range(X.shape[1]):
        Z[i, j] = (X[i, j] - mu[j]) / sd[j]

# 向量化写法：一次表达式完成
Z = (X - mu) / sd       # (1000,5) 与 (5,) 广播

这里 X - mu 触发广播：mu 形状为 (5,)，右对齐后与 X 的最后一维匹配，于是被拉伸到每一行；除法同理。整个过程中没有 Python 层的循环，计算全部落在 ufunc 的 C 循环里。若把 mu 误写成形状 (1000,)，则会因最右维 5 与 1000 不兼容而报错，这正是广播规则在帮你提前发现维度语义错误。

再看一个带条件的例子：把数组中所有负值截断为 0。循环写法需要逐个判断，而向量化可以直接用 np.maximum(a, 0) 或 np.where(a < 0, 0, a)。前者是单个 ufunc，通常更快；后者会先算掩码再选择，语义更直观。类似地，np.clip(a, 0, 1) 可一次性完成上下界截断。

四、改写时的常见陷阱

类型提升与溢出。 ufunc 会按输入 dtype 决定计算类型。若两个 int8 数组相加，结果仍是 int8，超过 127 就会溢出回绕。需要更大范围时应显式转换，例如 a.astype(np.int64) + b，或用 dtype= 参数指定。

临时数组与内存。 链式表达式 (X - mu) / sd 会先产生 X - mu 的临时数组，再产生结果。对超大数组，可用 out= 复用缓冲区，或分块处理，避免峰值内存翻倍。

广播的“隐式复制”错觉。 广播本身不复制数据，但结果数组是真实分配的。若把形状 (1000, 1) 与 (1, 1000) 相加，结果就是 (1000, 1000)，内存占用是百万级元素，需提前评估。

不要为了向量化而牺牲可读性。 当逻辑涉及跨元素依赖（如累加、滑动窗口、状态机）时，强行向量化往往需要复杂的索引技巧，反而易错。此时可优先考虑 np.cumsum、np.add.reduceat、np.lib.stride_tricks.sliding_window_view 等专用工具，或保留局部循环。

五、一个可复用的改写检查清单

掌握 ufunc 与广播之后，向量化改写就有了明确的判断依据：先看运算能否落到 ufunc 上，再看形状能否通过广播对齐，最后检查 dtype 与内存。把这三步走顺，绝大多数逐元素循环都能被替换为简洁、可读且更快的数组表达式。