如何扩展NumPy数组以实现与任意维度数组的广播运算？

当然有，你完全不需要手动判断输入数组的维度，利用NumPy的广播机制和轴扩展操作就能实现通用的相加逻辑，不用写一堆if-elif分支。

优化后的实现代码 #

这里有两种简洁的写法：

方法1：利用元组索引添加新轴

import numpy as np

def myfunction(input_array):
    array_to_add = np.array([1, 2, 3])
    # 计算需要补充的新轴数量：输入数组维度 - 1（因为array_to_add是1维）
    num_extra_axes = input_array.ndim - 1
    # 通过索引语法给array_to_add添加对应数量的长度为1的新轴
    expanded_array = array_to_add[(slice(None),) + (None,) * num_extra_axes]
    return input_array + expanded_array

方法2：使用np.expand_dims显式扩展轴

import numpy as np

def myfunction(input_array):
    array_to_add = np.array([1, 2, 3])
    # 在第1到第n-1位的位置添加新轴（n为输入数组的维度数）
    axes_to_expand = tuple(range(1, input_array.ndim))
    expanded_array = np.expand_dims(array_to_add, axes_to_expand)
    return input_array + expanded_array

原理说明 #

NumPy的广播规则允许维度匹配的数组进行元素级运算：

• 我们需要让array_to_add的形状和输入数组的维度对齐：输入数组第一个维度固定为3，和array_to_add的维度对应，后面的维度通过添加长度为1的新轴来匹配；
• 广播会自动把这些长度为1的轴扩展成输入数组对应维度的长度，完成逐元素相加。

比如：

• 输入为(3,)时，expanded_array仍为(3,)，直接相加；
• 输入为(3,5)时，expanded_array变为(3,1)，广播后与(3,5)完成相加；
• 输入为(3,5,4)时，expanded_array变为(3,1,1)，广播后完成对应位置的运算。

内容的提问来源于stack exchange，提问作者Alex V.