如何使用numpy.vectorize或广播实现指定规则的多维数组自定义运算

结论 #

两种方案都可实现，优先选择numpy.vectorize加signature参数的方案，写法更简洁，适配自定义函数的场景；如果你的自定义函数f本身支持向量化批量运算，用广播+原生numpy操作的效率更高。

你之前调试失败大概率是signature参数写法错误，或是没有对齐两个输入数组的广播维度。

可直接运行的vectorize实现方案 #

该方案完全匹配你要求的c[i, j] = f(a[i, j, :], b[i])计算逻辑：

import numpy as np

# 示例构造输入数组
a = np.random.rand(9, 3, 11)
b = np.random.rand(9)

# 你的自定义函数f，这里用求和乘b0做示例
def f(a0, b0):
    # 传入的a0长度固定为11，b0是对应i位置的标量
    # 注：你描述中写的len(b0) == 9应为笔误，按计算逻辑b0是单个值，若确实需要长度9的入参可自行调整逻辑
    return np.sum(a0) * b0

# 定义向量化函数，signature参数含义：第一个入参是长度为n的一维数组，第二个是标量，输出为标量
f_vec = np.vectorize(f, signature='(n),()->()')

# 扩展b的维度匹配a的前两维广播要求，直接计算
c = f_vec(a, b[:, np.newaxis])

# 验证输出形状为预期的(9,3)
print(c.shape)

广播适配方案 #

如果你的f支持批量输入，可以用维度调整+批量运算实现，性能略高于vectorize：

# 把a展平为(9*3, 11)的二维数组
a_flat = a.reshape(-1, 11)
# 把b重复3次，匹配a的第一维长度
b_repeat = np.repeat(b, 3)

# 批量计算后reshape回(9,3)
c = np.array([f(a_row, b_val) for a_row, b_val in zip(a_flat, b_repeat)]).reshape(9, 3)

如果f本身可以用numpy原生操作实现向量化，还可以进一步去掉列表推导，直接做批量运算，性能会有数量级提升。

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