如何通用化Numpy数组广播式扩展以实现多维矩阵与数组相乘？

当需要让一维数组N和多维数组M在指定轴上进行广播相乘时，手动添加None的方式扩展性太差——换个维度数或换个对应轴就得重新写索引。以下是两种通用化的实现方式：

方法一：动态构造索引元组 #

通过遍历M的维度索引，动态生成N的扩展索引，无需硬编码None的数量：

import numpy as np

# 示例配置：可根据实际场景修改
M = np.random.randn(2, 3, 4, 5)  # M的形状为(2,3,4,5)
axis = 2  # 指定N对应M的第2轴（索引从0开始）
N = np.random.randn(4)  # N的长度与M第2轴的维度一致

# 生成动态扩展索引：目标轴保留N的原维度，其余轴扩展为1
expand_indices = tuple(
    slice(None) if i == axis else None
    for i in range(M.ndim)
)

# 扩展N的维度并执行相乘
N_expanded = N[expand_indices]
result = M * N_expanded

逻辑说明 #

• M.ndim获取M的总维度数，遍历每个轴的索引i；
• 当i等于目标轴axis时，用slice(None)保留N的原维度（等价于:）；
• 其他轴用None添加长度为1的维度，最终生成的索引元组会自动匹配M的形状，实现广播。

方法二：构造目标形状后reshape #

先创建与M维度数一致的形状模板，仅在目标轴位置保留N的长度，其余设为1，再通过reshape扩展N：

import numpy as np

# 示例配置
M = np.random.randn(5, 6)  # M的形状为(5,6)
axis = 1  # 指定N对应M的第1轴
N = np.random.randn(6)

# 构造目标形状：与M维度数相同，目标轴为N的长度，其余为1
target_shape = [1] * M.ndim
target_shape[axis] = len(N)

# 扩展N并相乘
N_expanded = N.reshape(target_shape)
result = M * N_expanded

逻辑说明 #

• [1] * M.ndim生成一个全为1的列表，长度等于M的维度数；
• 将目标轴位置的元素替换为N的长度（需保证len(N)等于M.shape[axis]）；
• 通过reshape将N转换为目标形状，即可与M进行广播相乘。

这两种方法都能适配任意维度的M和任意指定轴的场景，无需手动调整索引或维度扩展的代码。

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