计算多对向量的点积的更好方法？

问题描述

我有两个数组 A 和 B 形状 (N,3)；我需要计算一个数组 C)，形状为 (N,1，使得 C[i]=A[i].B[i],i in range(N)。现在我正在做

import numpy as np
...
C = np.array([np.dot(a,b) for a,b in zip(A,B)]

有没有更快的方法？

解决方法

这是一种解决方案。不知道有没有更好的。

>>> A = np.random.randint(0,10,(3,3))
>>> A
array([[5,7,8],[6,1,2]])
>>> B = np.random.randint(0,3))
>>> B
array([[3,[2,6,0],[1,1]])
>>>
>>> np.dot(A[0],B[0])
86
>>> np.dot(A[1],B[1])
54
>>> np.dot(A[2],B[2])
14
>>>
>>> np.sum(A * B,axis=1)
array([86,54,14])

编辑：有一个更好的解决方案。使用 np.einsum

>>> np.einsum("ij,ij->i",A,B)
array([86,14])

我以前从未真正使用过 einsum，但一直都知道它。很酷的功能。使用爱因斯坦符号。