How can I optimize this pairwise Euclidean distance search algorithm?

Question

• x and y are two-dimensional arrays with dimensions (AxN) and (BxN), i.e. they have the same number of columns.

• I need to get a matrix of Euclidean distances between each pair of rows from x and y.

I have already written four different methods that give equally good results. But the first three methods are not fast enough, and the fourth method consumes too much memory.

1:

from scipy.spatial import distance

def euclidean_distance(x, y):
    return distance.cdist(x, y)

2:

import numpy as np
import math

def euclidean_distance(x, y):
    res = []
    one_res = []
    for i in range(len(x)):
        for j in range(len(y)):
            one_res.append(math.sqrt(sum(x[i] ** 2) + sum(y[j] ** 2) - 2 * 
               np.dot(x[i], np.transpose(y[j]))))
        res.append(one_res)
        one_res = []
    return res

3:

import numpy as np

def euclidean_distance(x, y):
    distances = np.zeros((x.T[0].shape[0], y.T[0].shape[0]))
    for x, y in zip(x.T, y.T):
        distances = np.subtract.outer(x, y) ** 2 + distances
    return np.sqrt(distances)

4:

import numpy as np


def euclidean_distance(x, y):
    x = np.expand_dims(x, axis=1)
    y = np.expand_dims(y, axis=0)
    return np.sqrt(((x - y) ** 2).sum(-1))

It looks like the second way could still be improved, and it also shows the best timing. Can you help me with the optimization, please?

Answer 1

Performance:

A way to use broadcasting but avoid allocating the 3d array is to separate the distance formula into $(x - y) ** 2 = x^2 + y^2 - 2xy$.

squared_matrix = np.sum(A ** 2,axis=1 )[:,np.newaxis] - 2 * (A@B.T) + np.sum(B ** 2,axis=1 )[np.newaxis,:]
distance_matrix = np.sqrt(squared_matrix)

Check:

np.allclose(distance_matrix,euclidean_distance(A,B))
>>> True

Answer 2

Ideas for 2:

• preallocate res
• don't use one_res, just a comprehensions