Better way to compute the dot product of many pairs of vectors?

Viewed 37

I have two arrays A and B shaped (N,3); I need to compute an array C), shaped (N,1 such that C[i]=A[i].B[i], i in range(N). Right now I'm doing

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

Is there a faster way?

1 Answers

Here's one solution. Not sure if there's something better.

>>> A = np.random.randint(0, 10, (3,3))
>>> A
array([[5, 7, 8],
       [6, 7, 8],
       [6, 1, 2]])
>>> B = np.random.randint(0, 10, (3,3))
>>> B
array([[3, 1, 8],
       [2, 6, 0],
       [1, 6, 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])

EDIT: there is a better solution. Using np.einsum

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

I've never actually used einsum before, but always knew about it. Very cool function. Uses Einstein notation.

Related