python - 根据索引计算从numpy数组中的所有点到单个点的距离

Question

假设一个二维数组如下:

arr = array([[1, 1, 1],
             [4, 5, 8],
             [2, 6, 9]])

如果给出了 point=array([1,1])，那么我想计算从 arr 的所有索引到点 (1,1) 的欧几里德距离。结果应该是

array([[1.41  , 1.        , 1.41],
       [1.    , 0.        , 1.  ],
       [1.41  , 1.        , 1.41]])

For 循环太慢而无法进行这些计算。有没有更快的方法来使用 numpy 或 scipy 来实现这一点？

谢谢!!!

Answer 1

方法 #1

您可以使用 scipy.ndimage.morphology.distance_transform_edt -

def distmat(a, index):
    mask = np.ones(a.shape, dtype=bool)
    mask[index[0],index[1]] = False
    return distance_transform_edt(mask)

方法 #2

另一个使用 NumPy 原生工具 -

def distmat_v2(a, index):
    i,j = np.indices(a.shape, sparse=True)
    return np.sqrt((i-index[0])**2 + (j-index[1])**2)

sample 运行 -

In [60]: a
Out[60]: 
array([[1, 1, 1],
       [4, 5, 8],
       [2, 6, 9]])

In [61]: distmat(a, index=[1,1])
Out[61]: 
array([[1.41421356, 1.        , 1.41421356],
       [1.        , 0.        , 1.        ],
       [1.41421356, 1.        , 1.41421356]])

In [62]: distmat_v2(a, index=[1,1])
Out[62]: 
array([[1.41421356, 1.        , 1.41421356],
       [1.        , 0.        , 1.        ],
       [1.41421356, 1.        , 1.41421356]])

基准测试

其他建议的解决方案:

# https://stackoverflow.com/a/61629292/3293881 @Ehsan
def norm_method(arr, point):
    point = np.asarray(point)
    return np.linalg.norm(np.indices(arr.shape, sparse=True)-point)

使用 benchit 包(很少有基准测试工具打包在一起；免责声明:我是它的作者)对提议的解决方案进行基准测试。

In [66]: import benchit

In [76]: funcs = [distmat, distmat_v2, norm_method]

In [77]: inputs = {n:(np.random.rand(n,n),[1,1]) for n in [3,10,50,100,500,1000,2000,5000]}

In [83]: T = benchit.timings(funcs, inputs, multivar=True, input_name='Length')

In [84]: In [33]: T.plot(logx=True, colormap='Dark2', savepath='plot.png')

因此， distmat_v2 似乎做得非常好，我们可以通过利用 numexpr 进一步改进它。

扩展到索引数组

我们可以扩展列出的解决方案以涵盖索引列表/数组 w.r.t. 的通用/更大的情况。我们需要在其余位置获得欧几里得距离，就像这样 -

def distmat_indices(a, indices):
    indices = np.atleast_2d(indices)
    mask = np.ones(a.shape, dtype=bool)
    mask[indices[:,0],indices[:,1]] = False
    return distance_transform_edt(mask)

def distmat_indices_v2(a, indices):
    indices = np.atleast_2d(indices)
    i,j = np.indices(a.shape, sparse=True)
    return np.sqrt(((i-indices[:,0])[...,None])**2 + (j-indices[:,1,None])**2).min(1)

sample 运行 -

In [143]: a = np.random.rand(4,5)

In [144]: distmat_indices(a, indices=[[2,2],[0,3]])
Out[144]: 
array([[2.82842712, 2.        , 1.        , 0.        , 1.        ],
       [2.23606798, 1.41421356, 1.        , 1.        , 1.41421356],
       [2.        , 1.        , 0.        , 1.        , 2.        ],
       [2.23606798, 1.41421356, 1.        , 1.41421356, 2.23606798]])