Python中带广播的逐元素多维矩阵乘法高效实现方法

你可以直接利用NumPy的广播机制来实现，不需要repeat这类会复制数据的操作，效率更高。核心思路是给两个数组添加合适的维度，让它们能通过广播对齐后逐元素相乘：

import numpy as np

a = np.arange(24).reshape((2,12))
b = np.arange(36).reshape((3,12))

# 给a增加第2个维度（变成(2,1,12)），给b增加第1个维度（变成(1,3,12)）
c = a[:, np.newaxis, :] * b[np.newaxis, :, :]

运行这段代码得到的结果和你原来的实现完全一致，但步骤更简洁，且内存效率更高——因为广播不会像repeat那样创建额外的数组副本，只是在计算时虚拟扩展维度。

原理说明 #

• a[:, np.newaxis, :] 将原形状(2,12)的数组转换为(2,1,12)
• b[np.newaxis, :, :] 将原形状(3,12)的数组转换为(1,3,12)
• 广播会自动将这两个数组扩展为(2,3,12)的形状（a的第2维度从1扩展到3，b的第1维度从1扩展到2），之后逐元素相乘就得到了你需要的结果。

性能对比 #

你的原方法中repeat会实际复制数据，当数组规模较大时，内存占用和计算时间都会显著增加；而基于广播的方法完全避免了不必要的数据复制，是NumPy处理这类维度扩展运算的标准高效方案。

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