Мне нравится этот, так как он работает для 1, 2 и 3D массивов
def e_dist(a, b, metric='euclidean'):
"""Distance calculation for 1D, 2D and 3D points using einsum
preprocessing :
use `_view_`, `_new_view_` or `_reshape_` with structured/recarrays
Parameters
----------
a, b : array like
Inputs, list, tuple, array in 1, 2 or 3D form
metric : string
euclidean ('e', 'eu'...), sqeuclidean ('s', 'sq'...),
Notes
-----
mini e_dist for 2d points array and a single point
>>> def e_2d(a, p):
diff = a - p[np.newaxis, :] # a and p are ndarrays
return np.sqrt(np.einsum('ij,ij->i', diff, diff))
See Also
--------
cartesian_dist : function
Produces pairs of x,y coordinates and the distance, without duplicates.
"""
a = np.asarray(a)
b = np.atleast_2d(b)
a_dim = a.ndim
b_dim = b.ndim
if a_dim == 1:
a = a.reshape(1, 1, a.shape[0])
if a_dim >= 2:
a = a.reshape(np.prod(a.shape[:-1]), 1, a.shape[-1])
if b_dim > 2:
b = b.reshape(np.prod(b.shape[:-1]), b.shape[-1])
diff = a - b
dist_arr = np.einsum('ijk,ijk->ij', diff, diff)
if metric[:1] == 'e':
dist_arr = np.sqrt(dist_arr)
dist_arr = np.squeeze(dist_arr)
return dist_arr
Выход
a = np.random.rand(3, 2)
b = np.random.rand(5, 2)
e_dist(a, b)
array([[0.62, 0.45, 0.88, 0.7 , 0.33],
[0.46, 0.57, 0.47, 0.25, 0.4 ],
[0.94, 0.68, 0.16, 0.35, 0.62]])
и для ...
a = np.random.rand(2, 3, 2)
b = np.random.rand(2, 5, 2)
e_dist(a, b)
array([[0.72, 0.39, 0.89, 0.25, 0.29, 0.41, 0.31, 0.6 , 0.8 , 0.39],
[0.25, 0.26, 0.69, 0.45, 0.45, 0.63, 0.57, 0.39, 0.25, 0.78],
[0.28, 0.27, 0.74, 0.47, 0.46, 0.69, 0.62, 0.45, 0.25, 0.84],
[0.43, 0.48, 0.47, 0.52, 0.54, 0.2 , 0.17, 0.21, 0.55, 0.38],
[0.15, 0.36, 0.54, 0.51, 0.52, 0.52, 0.48, 0.23, 0.23, 0.69],
[0.87, 0.76, 0.78, 0.67, 0.71, 0.26, 0.28, 0.64, 0.99, 0.06]])
И для сравнения вариантов numpy и scipy для 2d
from scipy.spatial import distance_matrix
distance_matrix(a, b)
array([[0.87, 0.89, 0.91, 0.07, 0.68, 0.95, 0.89],
[0.78, 0.52, 0.28, 0.85, 0.22, 0.65, 0.51],
[0.75, 0.46, 0.51, 1.08, 0.5 , 0.57, 0.45]])
e_dist(a, b)
array([[0.87, 0.89, 0.91, 0.07, 0.68, 0.95, 0.89],
[0.78, 0.52, 0.28, 0.85, 0.22, 0.65, 0.51],
[0.75, 0.46, 0.51, 1.08, 0.5 , 0.57, 0.45]])
Таким образом, существует множество вариантов, в зависимости от размера массивов, с которыми вы работаете.