如何用NumPy实现三维数组与二维数组的指定逐元素乘法？

核心实现代码 #

直接借助NumPy的广播机制，给二维数组a扩展一个维度后和三维数组b做逐元素乘法，全程无Python循环：

import numpy as np

# 两种等价写法选其一即可
# 写法1：用索引扩展维度
c = a[..., np.newaxis] * b

# 写法2：用expand_dims更直观
c = np.expand_dims(a, axis=-1) * b

原理简单说明 #

• 你的a是4×4的二维数组，每个a[i,j]是单个标量；b是4×4×2的三维数组，每个b[i,j]是长度为2的向量。
• NumPy要求运算的数组维度匹配才能逐元素计算，我们给a的最后加一个长度为1的维度（变成4×4×1），这样a的每个标量就会被自动“复制”成和b[i,j]长度一致的向量，完成对应位置的乘法。

验证示例效果 #

用你提供的测试数据运行，完全符合预期：

dist = np.array([[0.2, 0.3], [0.4,0.5], [0.6, 0.7], [0.8,0.9]])
a = np.array([[0.1, 1, 2, 3],[1,0.1,4,5],[2,4,0.1,6],[3,5,6,0.1]])
b = np.array([[[0,0], [-0.2,-0.2], [-0.4,-0.4], [-0.6,-0.6]], [[0.2,0.2], [0,0], [-0.2,-0.2], [-0.4,-0.4]], [[0.4,0.4], [0.2,0.2], [0,0], [-0.2,-0.2]],[[0.6,0.6], [0.4,0.4], [0.2,0.2], [0,0]]])

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

print(c[0,0])  # 输出 [0. 0.]，和预期一致
print(c[1,2])  # 输出 [-0.8 -0.8]，和预期一致

效率优势 #

这种实现完全依赖NumPy的底层C语言运算逻辑，没有Python层面的循环开销，速度远快于手动写for循环，非常适合需要数千次迭代的场景。

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