如何以NumPy风格实现不同形状矩阵的向量减法（替代for looping）

原始定义与循环实现 #

首先给出原始矩阵定义：

import numpy as np

a = np.array([[1,2], [3,4], [5,6], [7,8]])  # 形状(4, 2)，类型int32
b = np.array([[9,10], [11,12], [13,14]])    # 形状(3, 2)，类型int32

原先通过循环实现的逻辑是：取b的每一行与a做减法，最终得到形状为(3, 4, 2)的结果c，循环代码如下：

c = []
for i in range(b.shape[0]):
    c.append(b[i] - a)
c = np.array(c)

向量化实现方案 #

利用NumPy的广播机制可以完全替代循环，借助底层C优化提升运算效率，以下是几种等价的实现方式：

方式1：使用np.newaxis扩展维度

给b增加一个中间维度（从(3,2)变为(3,1,2)），让其与a的(4,2)形状兼容，触发广播完成逐行减法：

c = b[:, np.newaxis, :] - a

方式2：使用np.expand_dims显式扩展维度

如果觉得np.newaxis不够直观，可以用np.expand_dims指定要扩展的轴：

c = np.expand_dims(b, axis=1) - a

方式3：使用None替代np.newaxis

在NumPy中None和np.newaxis功能完全等价，写法更简洁：

c = b[:, None, :] - a

原理说明 #

NumPy的广播机制会自动将低维数组的维度扩展，使其与高维数组的形状匹配：

• 扩展后的b形状为(3,1,2)，a的形状为(4,2)
• 广播后两者都会被扩展为(3,4,2)，随后执行元素级减法运算，最终得到与循环实现完全一致的结果。

内容的提问来源于stack exchange，提问作者Muhammad Ikhwan Perwira