用numpy计算距离矩阵

Question

我正在尝试使用numpy来计算“距离矩阵”矩阵。更清楚地说，我有一个表示二维网格中位置的矩阵：

array([[[0, 0],
    [1, 0],
    [2, 0],
    [3, 0]],

   [[0, 1],
    [1, 1],
    [2, 1],
    [3, 1]],

   [[0, 2],
    [1, 2],
    [2, 2],
    [3, 2]],

   [[0, 3],
    [1, 3],
    [2, 3],
    [3, 3]],

   [[0, 4],
    [1, 4],
    [2, 4],
    [3, 4]]])

我可以遍历位置并计算目标位置和位置矩阵的每个位置之间差异的范数，如下所示：

pos_goal = np.array([1,2])
dist_matrix = np.zeros(l_arr.shape[:2])
for i, line in enumerate(l_arr):
    for j, pos in enumerate(line):
        dist_matrix[i,j] = np.linalg.norm(pos - pos_goal) 

dist_matrix

结果：

array([[ 2.23606798,  2.        ,  2.23606798,  2.82842712],
   [ 1.41421356,  1.        ,  1.41421356,  2.23606798],
   [ 1.        ,  0.        ,  1.        ,  2.        ],
   [ 1.41421356,  1.        ,  1.41421356,  2.23606798],
   [ 2.23606798,  2.        ,  2.23606798,  2.82842712]])

有没有更好的方法来做到这一点（没有2个循环）？

Answer 1

使用scipy.spatial.distance.cdist。它需要2D输入，因此您可以执行以下操作：

from scipy.spatial import distance
dist_matrix = distance.cdist(l_arr.reshape(-1, 2), [pos_goal]).reshape(l_arr.shape[:2])

这非常简洁，对于大型阵列而言，比基于循环或广播的手动方法更快。

Answer 2

np.linalg.norm函数采用axis参数，因此您需要：

In [6]: np.linalg.norm(l_arr - pos_goal, axis=2)
Out[6]:
array([[ 2.23606798,  2.        ,  2.23606798,  2.82842712],
       [ 1.41421356,  1.        ,  1.41421356,  2.23606798],
       [ 1.        ,  0.        ,  1.        ,  2.        ],
       [ 1.41421356,  1.        ,  1.41421356,  2.23606798],
       [ 2.23606798,  2.        ,  2.23606798,  2.82842712]])

你可以使用-1作为“最后一个”轴：

In [7]: np.linalg.norm(l_arr - pos_goal, axis=-1)
Out[7]:
array([[ 2.23606798,  2.        ,  2.23606798,  2.82842712],
       [ 1.41421356,  1.        ,  1.41421356,  2.23606798],
       [ 1.        ,  0.        ,  1.        ,  2.        ],
       [ 1.41421356,  1.        ,  1.41421356,  2.23606798],
       [ 2.23606798,  2.        ,  2.23606798,  2.82842712]])

注意，我使用数组广播来获得差异：

In [11]: l_arr - pos_goal
Out[11]:
array([[[-1, -2],
        [ 0, -2],
        [ 1, -2],
        [ 2, -2]],

       [[-1, -1],
        [ 0, -1],
        [ 1, -1],
        [ 2, -1]],

       [[-1,  0],
        [ 0,  0],
        [ 1,  0],
        [ 2,  0]],

       [[-1,  1],
        [ 0,  1],
        [ 1,  1],
        [ 2,  1]],

       [[-1,  2],
        [ 0,  2],
        [ 1,  2],
        [ 2,  2]]])

通常，learning how to use broadcasting结合内置的numpy / scipy矢量化函数是实现显着速度提升的方法。