python - NumPy 广播 : Calculating sum of squared differences between two arrays

Question

我有以下代码。在 Python 中它需要永远。必须有一种方法可以将这种计算转化为广播......

def euclidean_square(a,b):
    squares = np.zeros((a.shape[0],b.shape[0]))
    for i in range(squares.shape[0]):
        for j in range(squares.shape[1]):
            diff = a[i,:] - b[j,:]
            sqr = diff**2.0
            squares[i,j] = np.sum(sqr)
    return squares

Answer 1

您可以使用 np.einsum在计算出 broadcasted way 中的差异之后, 像这样 -

ab = a[:,None,:] - b
out = np.einsum('ijk,ijk->ij',ab,ab)

或者使用scipy's cdist将其可选的度量参数设置为 'sqeuclidean' 为我们提供问题所需的平方欧氏距离，就像这样 -

from scipy.spatial.distance import cdist
out = cdist(a,b,'sqeuclidean')