Calculate nearest neighbors using a nxn matrix-CodePudding

Supposed I have nxn T random matrix, how can I find nearest neighbors element of the matrix?

enter image description here

CodePudding user response：

According to your image description, my usual way is to generate an index matrix, calculate the Euclidean distance between each element and a given position, and then judge whether the distance is 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])

The advantage of this method are:

it is usually not slow, unless your matrix is very large, but there are still workarounds.
you needn't to consider some boundary conditions. For example, the target position is at the edge of the matrix.
it has strong expansibility. For example, if you redefine the neighbor as an element with a distance of less than or equal to 2, or an element in a circle around itself, you can achieve your expectation with only a little modification.

For a large matrix, considering the boundary conditions, you can implement it through such a function:

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