如何优雅替代NumPy单轴循环实现指定数组运算

你可以利用NumPy的广播机制，通过给数组b添加一个新的最后维度，直接和a做元素级乘法，完全替代循环操作，代码更简洁且符合NumPy惯用风格：

import numpy as np

# 核心实现
c = a * b[..., np.newaxis]

原理说明 #

你的原代码是循环a的最后一维，将每个切片a[...,i]与b相乘后赋值给c的对应位置。而通过b[..., np.newaxis]，我们给b增加了一个长度为1的末尾维度，让b的形状从b.shape变为b.shape + (1,)。

由于a.shape[:-1]和b.shape本身具备可广播性，扩展后的b与a的所有维度都能匹配：a的形状是(..., N)，扩展后的b形状是(..., 1)，NumPy会自动将b的最后一维广播到长度N，完成逐元素相乘，最终得到的结果和原循环代码完全一致。

等价写法 #

你也可以用np.expand_dims来显式扩展维度，效果相同：

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

验证正确性 #

可以用以下代码验证两种方式的结果一致：

# 示例测试数组
a = np.random.rand(2, 3, 4)
b = np.random.rand(2, 3)

# 原循环实现
c_old = np.empty(np.broadcast_shapes(a.shape[:-1], b.shape) + (a.shape[-1],), a.dtype)
for i in range(a.shape[-1]):
    c_old[..., i] = a[..., i] * b

# 新实现
c_new = a * b[..., np.newaxis]

# 断言结果一致（考虑浮点精度误差用allclose）
assert np.allclose(c_old, c_new)

这种向量化方式不仅代码更简洁，还能充分利用NumPy的底层优化，比Python循环的执行效率高得多，尤其是当a的最后维度长度很大时，性能提升会非常明显。

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