如何向量化实现矩阵行与另一矩阵的整体相加操作？

现有r×k的矩阵A和m×k的矩阵B，需要将A的每一行A[i,:]与完整的B相加，最终得到r个独立的m×k矩阵（等价于生成r×m×k的3D矩阵）。

原实现通过Python循环完成逻辑，虽能得到预期输出，但运行速度较慢：

import numpy as np

v = np.array([[1,2],[3,4]])
w = np.zeros((5,2))

def mysum(a,b):
    return np.array([a[i]+b for i in range(a.shape[0])])

mysum(v,w)

预期输出：

[array([[1., 2.],
        [1., 2.],
        [1., 2.],
        [1., 2.],
        [1., 2.]]),
 array([[3., 4.],
        [3., 4.],
        [3., 4.],
        [3., 4.],
        [3., 4.]])]

高效向量化实现方案 #

利用numpy的广播机制，通过调整数组维度实现无Python循环的元素级相加，这是numpy性能优化的核心手段。

实现代码 #

import numpy as np

v = np.array([[1,2],[3,4]])
w = np.zeros((5,2))

# 扩展v的维度为(r,1,k)，与w的(m,k)维度自动广播匹配
result_3d = v[:, np.newaxis, :] + w

# 若需要和原输出格式一致（r个m×k矩阵的列表），直接拆分3D数组即可
result_list = list(result_3d)

等价写法（用reshape调整维度） #

# 用reshape替代np.newaxis，效果完全一致
result_3d = v.reshape(v.shape[0], 1, v.shape[1]) + w

原理说明 #

• 原矩阵A是(r,k)结构，通过[:, np.newaxis, :]扩展为(r,1,k)后，numpy会自动将这个新增的维度广播到与B的m维度匹配，最终每个A[i,:]会和B的所有行对应相加，生成(r,m,k)的3D数组。
• 广播操作由numpy底层C代码实现，完全规避了Python层面的循环开销，在大规模数据场景下性能提升极为明显。

性能优势 #

当处理r、m较大的矩阵时（比如r=1000，m=1000，k=100），广播实现的速度比原循环快几十至上百倍，且数据量越大，性能差距越显著。

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