使用 nxn 矩阵计算最近邻

Calculate nearest neighbors using a nxn matrix

假设我有 nxn T 个随机矩阵,如何找到矩阵的最近邻元素?

根据你的图片描述,我通常的做法是生成索引矩阵,计算每个元素到给定位置的欧氏距离,然后判断距离是否为1:

>>> m = np.random.rand(5, 5)
>>> m
array([[0.15993681, 0.34212   , 0.27103587, 0.47880503, 0.92797347],
       [0.56698489, 0.16124083, 0.63037649, 0.79914445, 0.55247301],
       [0.05826773, 0.26793143, 0.62384919, 0.8248818 , 0.75127237],
       [0.38831836, 0.51175059, 0.70764241, 0.08915257, 0.63465847],
       [0.17459285, 0.74014131, 0.42850117, 0.54282051, 0.61795445]])
>>> ii, jj = np.indices(m.shape)
>>> i, j = 2, 3
>>> mask = np.hypot(ii - i, jj - j) == 1
>>> mask
array([[False, False, False, False, False],
       [False, False, False,  True, False],
       [False, False,  True, False,  True],
       [False, False, False,  True, False],
       [False, False, False, False, False]])
>>> m[mask]
array([0.79914445, 0.62384919, 0.75127237, 0.08915257])

这种方法的优点是:

  • 一般不会慢,除非你的矩阵很大,不过还是有变通办法的。
  • 你不需要考虑一些边界条件。比如目标位置在矩阵的边缘。
  • 扩展性强。比如你把neighbor重新定义为一个距离小于等于2的元素,或者一个围绕自己的圆圈内的元素,只要稍微修改一下就可以达到你的预期。

对于大矩阵,考虑到边界条件,可以通过这样的函数实现:

>>> def neighbor_indices(arr, pos, dist_func=np.hypot, dist=1):
...     start = [max(p - dist, 0) for p in pos]
...     shape = [s - p for s, p in zip(arr.shape, start)]
...     indices = np.indices(shape)
...     ii, jj = [i + s for i, s in zip(indices, start)]
...     i, j = pos
...     mask = dist_func(ii - i, jj - j) == dist
...     return ii[mask], jj[mask]
...
>>> neighbor_indices(m, (2, 3))
(array([1, 2, 2, 3]), array([3, 2, 4, 3]))
>>> neighbor_indices(m, (3, 4))
(array([2, 3, 4]), array([4, 3, 4]))
>>> neighbor_indices(m, (-1, -1))
(array([], dtype=int32), array([], dtype=int32))
>>> neighbor_indices(m, (5, 5))
(array([], dtype=int32), array([], dtype=int32))
>>> neighbor_indices(m, (2, 3), dist=2)
(array([0, 2, 4]), array([3, 1, 3]))
>>> neighbor_indices(m, (2, 3), lambda i, j: np.maximum(np.abs(i), np.abs(j)))
(array([1, 1, 1, 2, 2, 3, 3, 3]), array([2, 3, 4, 2, 4, 2, 3, 4]))