numpy - 计算沿第一轴的两个数组之间的“元素方式”差异答案

【问题标题】：numpy - Computing "element-wise" difference between two arrays along first axisnumpy - 计算沿第一轴的两个数组之间的“元素方式”差异
【发布时间】：2016-08-31 13:38:06
【问题描述】：

假设我有两个数组A 和B，维度分别为(n1,m1,m2) 和(n2,m1,m2)。我想计算矩阵C，尺寸为(n1,n2)，使得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有类似的方法吗？

谢谢。

【问题讨论】：

标签： python arrays performance numpy vectorization

【解决方案1】：

您可以使用scipy's cdist，这对于在reshaping 输入数组到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)

【讨论】：