bordeo - 8 months ago 54

Python Question

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

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')
```