NumPy中计算距离矩阵的有效方法

hsgswve4  于 12个月前  发布在  其他
关注(0)|答案(3)|浏览(138)

我有这个示例数组:

In [38]: arr
Out[38]: array([  0,  44, 121, 154, 191])

字符串
上面只是一个例子,而我的实际数组大小是相当大的。那么,什么是一个有效的方法来计算距离矩阵?
结果应该是:

In [41]: res
Out[41]: 
array([[   0,   44,  121,  154,  191],
       [ -44,    0,   77,  110,  147],
       [-121,  -77,    0,   33,   70],
       [-154, -110,  -33,    0,   37],
       [-191, -147,  -70,  -37,    0]])


我写了一个基于for循环的实现,它太慢了。出于效率的原因,可以将它向量化吗?

kxkpmulp

kxkpmulp1#

可以使用broadcasting

from numpy import array

arr = array([  0,  44, 121, 154, 191])
arrM = arr.reshape(1, len(arr))
res = arrM - arrM.T

字符串

zzlelutf

zzlelutf2#

subtract.outer,它有效地执行两个数组之间的广播减法。
将ufunc op应用于所有对(a,B),其中a在A中,B在B中。
设M = A.ndim,N = B. ndim,则op.outer(A, B)的结果C是一个维数为M + N的数组,使得:

C[i_0, ..., i_{M-1}, j_0, ..., j_{N-1}] = 
     op(A[i_0, ..., i_{M-1}],B[j_0, ..., j_{N-1}])

个字符
或者说,

arr - arr[:, None] # essentially the same thing as above

array([[   0,   44,  121,  154,  191],
       [ -44,    0,   77,  110,  147],
       [-121,  -77,    0,   33,   70],
       [-154, -110,  -33,    0,   37],
       [-191, -147,  -70,  -37,    0]])

4ngedf3f

4ngedf3f3#

def dist_calc(
    arr1: np.ndarray[tuple[int, int], np.dtype[int | float]],
    arr2: np.ndarray[tuple[int, int], np.dtype[int | float]],
) -> np.ndarray[tuple[int, int], np.dtype[float]]:
    """
    calculates euclidean distances between all points in two k-dimensional arrays
    'arr1' and 'arr2'
        :parameter
            - arr1: N x k array
            - arr2: M x k array
        :return
            - dist: M x N array with pairwise distances
    """
    norm_1 = np.sum(arr1 * arr1, axis=1).reshape(1, -1)
    norm_2 = np.sum(arr2 * arr2, axis=1).reshape(-1, 1)

    dist = (norm_1 + norm_2) - 2.0 * np.dot(arr2, arr1.T)
    # necessary due to limited numerical accuracy
    dist[dist < 1.0e-11] = 0.0

    return np.sqrt(dist)

res = dist_calc(arr, arr)

字符串
这将是一个更一般的答案,其中数组不必是一维的。

相关问题