使用numpy计算距离矩阵

时间:2017-12-20 00:20:57

标签: python arrays numpy matrix

我正在尝试使用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个循环)?

2 个答案:

答案 0 :(得分: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矢量化函数相结合,是实现大幅提升速度的方法。

答案 1 :(得分:2)

使用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])

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

相关问题