如何优雅地在不同轴上实现NumPy数组相乘？

输入数组 #

先定义问题中的两个输入数组：

import numpy as np

a = [
     [[-1, 0, 1], 
      [-2, 0, 2],
      [-3, 0, 3]],

     [[-4, 0, 4],
      [-5, 0, 5],
      [-6, 0, 6]]
] 
b = [
     [1, 3, 5],
     [7, 11, 13]
]

需求分析 #

观察目标输出c的规律：c[i][j][k] = a[i][j][k] * b[i][j]，即b中每个位置的标量，需要与a中对应位置的长度为3的子数组做逐元素乘法。

实现方案 #

利用NumPy的广播机制就能优雅解决，无需编写循环。核心是给b增加一个维度，让它的形状从(2, 3)适配a的(2, 3, 3)：

# 转换为NumPy数组
a_np = np.array(a)
b_np = np.array(b)

# 扩展b的维度后执行逐元素乘法
c_np = a_np * b_np[..., np.newaxis]
# 也可以用reshape实现维度扩展：b_np.reshape(2, 3, 1)

# 转换为列表格式查看结果
print(c_np.tolist())

运行后得到的结果完全符合目标输出c。

原理说明 #

NumPy广播会自动扩展维度匹配的数组：当b_np被扩展为(2, 3, 1)的形状后，它的第三个维度会被广播至长度3，与a_np的第三个维度对齐，从而实现对应位置的子数组与标量的逐元素相乘。

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