如何对两个numpy数组的逐行外加法运算进行向量化以去除for循环

向量化实现方案 #

你可以直接用numpy的广播机制实现无循环运算，核心是给两个数组补充维度触发自动广播，代码如下：

import numpy as np

A = np.array([[ 0,  1,  2,  3,  4],
              [ 5,  6,  7,  8,  9],
              [10, 11, 12, 13, 14]])

B = np.array([[1., 1., 1., 1., 1., 1.],
              [2., 2., 2., 2., 2., 2.],
              [3., 3., 3., 3., 3., 3.]])

# 核心代码：补充维度后直接相加
C = A[:, :, None] + B[:, None, :]
# 也可以写为 C = A[..., np.newaxis] + B[:, np.newaxis, :]，None和np.newaxis功能完全等价

上述代码输出的C和你原有for循环得到的结果完全一致，可以通过np.allclose(C, 原循环生成的C)验证相等性。

逻辑说明 #

你原有循环中对每一行的处理逻辑是：将形状为(k,)的A[r]拓展为(k, 1)，和形状为(n,)的B[r]相加，通过广播得到(k, n)的单页结果。
向量化的逻辑就是把这个维度拓展操作批量应用到整个数组：

• 给形状为(m, k)的A补充最后一维，变成(m, k, 1)
• 给形状为(m, n)的B补充中间维度，变成(m, 1, n)

两个数组相加时，numpy会自动在长度为1的维度上做广播，直接生成形状为(m, k, n)的结果，所有运算都在numpy底层执行，比Python级别的for循环性能高很多，当m的数值很大时优势尤其明显。

内容的提问来源于stack exchange，提问作者PJORR