python - numpy - 计算沿第一个轴的两个数组之间的 "element-wise"差异

Question

假设我有两个数组 A 和 B，其维度为 (n1,m1,m2) 和 (n2,m1, m2) 分别。我想计算尺寸为 (n1,n2) 的矩阵 C，使得 C[i,j] = sum((A[i,:,: ] - B[j,:,:])^2)。这是我到目前为止所拥有的:

import numpy as np
A = np.array(range(1,13)).reshape(3,2,2)
B = np.array(range(1,9)).reshape(2,2,2)
C = np.zeros(shape=(A.shape[0], B.shape[0]) )
for i in range(A.shape[0]):
    for j in range(B.shape[0]):
        C[i,j] = np.sum(np.square(A[i,:,:] - B[j,:,:]))
C

最有效的方法是什么？在 R 中，我会使用矢量化方法，例如 outer。 Python 有类似的方法吗？

谢谢。

Answer 1

您可以使用scipy's cdist ，在将输入数组 reshape 为2D之后，这对于此类计算非常有效，如下所示 -

from scipy.spatial.distance import cdist

C = cdist(A.reshape(A.shape[0],-1),B.reshape(B.shape[0],-1),'sqeuclidean')

现在，上述方法必须具有内存效率，因此在处理大数据时是更好的方法。对于小型输入数组，还可以使用 np.einsum和杠杆NumPy broadcasting ，就像这样 -

diffs = A[:,None]-B
C = np.einsum('ijkl,ijkl->ij',diffs,diffs)