NumPy Broadcasting:计算两个数组之间的平方差之和

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在计算a中的差异后使用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')