我想添加一个一维数组到二维数组沿着二维数组的第二个维度使用以下代码中的逻辑。
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)
这个列表理解非常慢。什么是最快的方法?(我猜是用数组计算)。
EDIT我尝试了以下方法,但运行时间最差
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]
3条答案
按热度按时间piok6c0g1#
由于我们将相同的值(1D数组的每个元素)添加到2D数组中,这相当于将每个轴的总和加上1D数组**乘以2D数组第二个轴的维度。我们还必须向1D数组添加一个新的轴,否则形状不匹配。
和
uoifb46i2#
如果你想编写一个高性能的Numba实现,你必须编写简单的循环,或者在某些情况下使用BLAS函数。
实施例1
编辑:示例2
与一些想法从弥敦道炉回答这可以修改为
snz8szmq3#
Change - Sum = np.array([(TwoDArray+element).sum(axis=1)for element in OneDArray]).T