如何高效无循环将2D NumPy数组A与1D数组B在第三维度做减法生成3D数组C

已知2D NumPy数组A（形状4×5）和1D数组B（元素为[1, 100]），直接执行A-B会因列数不匹配报错。需要高效生成3D数组C，其中：

• C[0, i, j] = A[i,j] - B[0]
• C[1,i,j] = A[i,j] - B[1]
  要求避免编写大量循环，利用NumPy的广播机制完成。

原代码如下：

import numpy as np
B=np.array([1, 100])
A=np.arange(4*5).reshape(4,5)
#C=A-B this will not work as A and B have different number of columns
#A=array([[ 0,  1,  2,  3,  4],
#        [ 5,  6,  7,  8,  9],
#        [10, 11, 12, 13, 14],
#        [15, 16, 17, 18, 19]])

解决方案 #

利用NumPy的广播机制，调整数组维度让运算合法：

import numpy as np

B = np.array([1, 100])
A = np.arange(4*5).reshape(4,5)

# 将B的维度从(2,)扩展为(2,1,1)，匹配A的广播要求
C = A - B[:, np.newaxis, np.newaxis]

# 验证结果
print("C的形状：", C.shape)  # 输出 (2, 4, 5)
print("\nC[0]（A各元素减1）：")
print(C[0])
print("\nC[1]（A各元素减100）：")
print(C[1])

原理说明 #

NumPy广播规则：数组在每个维度上要么长度相同，要么某一方维度长度为1。

• 原A形状为(4,5)，B形状为(2,)
• 将B扩展为(2,1,1)后，A会被自动广播为(1,4,5)，两者运算时会广播为统一的(2,4,5)形状，正好生成符合要求的3D数组C。

也可以用np.expand_dims替代[:, np.newaxis, np.newaxis]，写法更直观：

C = A - np.expand_dims(B, axis=(1, 2))

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