Skip to content

Efficient Euclidean Distance Calculation - Numpy Einsum

This script uses numpy's einsum function to calculate the euclidean distance. Resources:

import numpy as np

def euclidean_distance_einsum(X, Y):
    """Efficiently calculates the euclidean distance
    between two vectors using Numpys einsum function.

    Parameters
    ----------
    X : array, (n_samples x d_dimensions)
    Y : array, (n_samples x d_dimensions)

    Returns
    -------
    D : array, (n_samples, n_samples)
    """
    XX = np.einsum('ij,ij->i', X, X)[:, np.newaxis]
    YY = np.einsum('ij,ij->i', Y, Y)
#    XY = 2 * np.einsum('ij,kj->ik', X, Y)
    XY = 2 * np.dot(X, Y.T)
    return  XX + YY - XY

An alternative way per stackoverflow would be to do it in one shot.

Sources