bordeo bordeo - 1 year ago 126
Python Question

NumPy Broadcasting: Calculating sum of squared differences between two arrays

I have the following code. It is taking forever in Python. There must be a way to translate this calculation into a broadcast...

def euclidean_square(a,b):
squares = np.zeros((a.shape[0],b.shape[0]))
for i in range(squares.shape[0]):
for j in range(squares.shape[1]):
diff = a[i,:] - b[j,:]
sqr = diff**2.0
squares[i,j] = np.sum(sqr)
return squares

Answer Source

You can use np.einsum after calculating the differences in a broadcasted way, like so -

ab = a[:,None,:] - b
out = np.einsum('ijk,ijk->ij',ab,ab)

Or use scipy's cdist with its optional metric argument set as 'sqeuclidean' to give us the squared euclidean distances as needed for our problem, like so -

from scipy.spatial.distance import cdist
out = cdist(a,b,'sqeuclidean')