计算numpy.ndarray内部乘法的有效方法

时间:2014-10-10 06:22:41

标签: python performance numpy matrix multidimensional-array

我有两个具有相同形状的矩阵a和b:

a = np.ndarray(shape=(3, 2), dtype=int)
b = np.ndarray(shape=(3, 2), dtype=int)

我希望它们的内部乘法如下:

    1 2
a = 4 5
    7 8

    9 8
b = 6 5
    3 2

我希望结果是这个1x2 ndarray:

[1x9+4x6+7x3   2x8+5x5+8x2]

就像我们正在计算两个矩阵的列的标量点积。

我现在正在做的是:

np.diag(np.dot(np.transpose(a), b))

但它效率不高,而且正在计算许多其他我不需要的东西。

我的矩阵比这些大得多,所以找到更有效的解决方案很重要。

1 个答案:

答案 0 :(得分:2)

您可以先进行简单乘法,然后在轴= 0上求和:

>>> a = np.array([[1, 2], [4, 5], [7, 8]])
>>> b = np.array([[9, 8], [6, 5], [3, 2]])
>>> (a * b).sum(axis=0)
array([54, 57])