如何在不新增轴的情况下高效广播(n,m,k)与(n,m)数组的乘法？

针对大数组内存不足的问题，你可以通过以下几种方式避免显式创建扩展轴的中间数组，同时实现相同的掩码效果：

• 原地修改原数组（最省内存）
  直接在原数组a上进行乘法操作，不需要额外创建新数组b，能大幅节省内存：

  a *= a_msk[..., None]
  

  这里[..., None]和[:,:,np.newaxis]作用一致，但写法更简洁，而且它创建的是原数组的视图而非副本，本身几乎不占用额外内存。原地修改后，a直接变为掩码后的结果。

• 使用np.einsum指定维度运算
  einsum可以通过维度标记直接实现不同维度数组的对应相乘，无需手动扩展轴，且内存效率较高：

  b = np.einsum('nmk,nm->nmk', a, a_msk)
  

  该方法会直接计算对应位置的乘积，避免了中间扩展数组的显式创建，适合需要保留原数组的场景。

• 利用广播机制的隐式扩展
  numpy的广播机制会自动处理维度匹配，你也可以通过reshape调整掩码数组的形状，结合原地操作同样能节省内存：

  a *= a_msk.reshape(a_msk.shape + (1,))
  

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