沿着3D阵列的轴的向量化欧式距离-Python

Question

如何在3D数组的轴上应用在2D数组（cdist）上运行的函数？

我尝试使用numpy.apply_along_axis，但是我需要在2D数组而不是1D数组上进行操作。我通过沿一个轴进行迭代获得了所需的结果，但我希望将其向量化：

from scipy import spatial
import numpy as np

a = np.random.randn(600).reshape(10, 20, 3)
distances = np.array([spatial.distance.cdist(a[i,:,:], a[i,:,:]) for i in range(a.shape[0])])

Answer 1

受this post的启发，我们可以以向量化的方式解决它。因此，遵循wiki contents包中的eucl_dist（免责声明：我是它的作者），我们可以利用matrix-multiplication和一些NumPy specific implementations，就像这样-

a_s = np.einsum('ijk,ijk->ij',a,a)
sq_dists = a_s[...,None]+a_s[...,None,:]-2*np.einsum('ijk,ilk->ijl',a,a)
dists = np.sqrt(sq_dists)

替代项：

• 我们可以使用np.matmul/@-operator on Python3.x来替换矩阵乘法部分。因此，np.einsum('ijk,ilk->ijl',a,a)可以替换为np.matmul(a,a.transpose(0,2,1))。