如何使用numpy向量化方法高效完成ndarray维度扩展与公式赋值

核心思路 #

利用numpy的广播机制实现向量化运算，完全避免显式for循环，运算效率远高于Python原生循环实现。

操作步骤 #

1. 计算新增维度的数值
   你给出的计算公式可以直接通过数组广播实现，不需要遍历i：

# 计算得到的new_channel形状为 (100, 50)
new_channel = A[:, :, 4] * B + 5
# 为新增数组扩展最后一维，方便后续拼接，扩展后形状为 (100, 50, 1)
new_channel = new_channel[..., np.newaxis]

2. 拼接数组得到最终结果
   使用np.concatenate沿最后一维拼接原数组和新生成的通道即可：

A_extended = np.concatenate([A, new_channel], axis=-1)

最终得到的A_extended形状就是你需要的[100, 50, 6]。

原理说明 #

A[:, :, 4]的形状为(100, 50)，权重数组B的形状为(50,)，二者做乘法运算时，numpy会自动将B沿第一个维度广播复制100次，刚好满足每个i位置对应乘B[i]的计算逻辑，完全匹配你给出的公式要求。

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