如何高效实现1D NumPy数组与2D数组指定维度的求和运算？

我希望按照下方代码的逻辑，将1D NumPy数组沿2D数组的第二维度进行相加运算：

import numpy as np
TwoDArray = np.random.randint(0, 10, size=(10000, 50))
OneDArray = np.random.randint(0, 10, size=(2000))
Sum = np.array([(TwoDArray+element).sum(axis=1) for element in OneDArray]).T
print(Sum.shape)
>> (10000, 2000)

这段列表推导式运行速度极慢，请问最快的实现方式是什么？（我推测可以用数组计算实现）

编辑补充
我尝试了以下方法，但运行速度更差：

Sum = np.sum(TwoDArray[:, np.newaxis, :] + OneDArray[np.newaxis, :, np.newaxis], axis=2)

我还尝试了Numba，但耗时与原方法相同：

@jit(nopython=True)
def compute_sum(TwoDArray, OneDArray):
    return [(TwoDArray+element).sum(axis=1) for element in OneDArray]

核心是通过数学推导简化计算逻辑，彻底避免不必要的大规模元素级运算：

原代码中，(TwoDArray + element).sum(axis=1) 等价于：
TwoDArray.sum(axis=1) + element * TwoDArray.shape[1]
原因很简单：给2D数组的每个元素加element后，每行的总和等于原行的总和加上element乘以该行的元素个数（也就是2D数组的列数）。

基于这个推导，用Numba广播直接计算即可，代码如下：

import numpy as np

TwoDArray = np.random.randint(0, 10, size=(10000, 50))
OneDArray = np.random.randint(0, 10, size=(2000))

# 预计算2D数组的所有行和，转成(10000,1)形状以便广播
row_sums = TwoDArray.sum(axis=1)[:, np.newaxis]
# 利用广播完成批量计算：行和数组 + 每个element*列数
Sum = row_sums + OneDArray * TwoDArray.shape[1]

print(Sum.shape)  # 输出 (10000, 2000)

为什么这个方法最快？ #

• 计算量骤降：原方法需要做 2000 * 10000 * 50 = 1e9 次加法运算，而新方法只需要 10000*50 次求和（预计算行和）加上 10000*2000 次加法，计算量不到原方法的1/20。
• 无冗余内存占用：不会创建超大的中间数组（比如之前尝试的三维广播会生成1e9个元素的数组，直接吃满内存带宽），仅用两个小维度数组通过广播计算。
• 完全利用NumPy底层优化：广播运算由NumPy的C级实现加速，比Python循环或Numba循环效率高得多。

为什么之前的尝试效果差？ #

• 三维广播方法：生成了(10000,2000,50)的巨型中间数组，内存开销极大，计算时频繁读写内存，速度被内存带宽限制。
• Numba版本：只是加速了Python循环，但没有改变“重复做1e9次加法”的核心逻辑，计算量和原方法一致，所以速度没提升。

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