查询numpy二维数组中一个元素的四个邻居的值

Question

I have a 2D array of 5*5 like this:

>>> np.random.seed(100)
>>> a = np.random.randint(0,100, (5,5))
>>> a
array([[ 8, 24, 67, 87, 79],
       [48, 10, 94, 52, 98],
       [53, 66, 98, 14, 34],
       [24, 15, 60, 58, 16],
       [ 9, 93, 86,  2, 27]])

if I have an initial position, is there any way to quickly and easily get the values of its four neighbors around it? The method I'm using now is a bit cumbersome:

Suppose the current position is [x, y] (if x=2, y=3 then the value in the array is 14,)，then the position above it is [x-1, y], the bottom is [x+1, y], the left side is [y-1, x], and the right side is [y+1, x]. I use the following four lines of code to get the values of neighbors.

curr_val = a[2,3]
up_val = a[2+1, 3]
bott_val = a[2-1, 3]
left_val = a[2, 3+1]
right_val = a[2, 3-1]

So my question is is there a more convenient function in numpy that can do this and even query the values of four neighbors at once?

Answer 1

This is not the shortest method, but a flexible way could be to use a mask and a convolution to build this mask.

The advantage is that you can use any mask easily, just change the kernel.

from scipy.signal import convolve2d

kernel = [[0,1,0],  # define points to pick around the target
          [1,0,1],
          [0,1,0]]
mask = np.zeros_like(a, dtype=bool) # build empty mask
mask[x,y] = True                    # set target(s)

# boolean indexing
a[convolve2d(mask, kernel, mode='same').astype(bool)]

output: array([52, 98, 34, 58])

Answer 2

The fastest way is this one taking usec to compute. Some times shortest is not the best. This one is very simple to understand and has no package dependencies.

This also works for edge-cases.

def neighbors(matrix: np.ndarray, x: int, y: int):
    x_len, y_len = matrix.shape
    nbr = []
    if x != 0:
        nbr.append(matrix[x-1][y])
    if y != 0:
        nbr.append(matrix[x][y-1])
    if x != x_len:
        nbr.append(matrix[x+1][y])
    if y != y_len:
        nbr.append(matrix[x][y+1])
    return nbr

Answer 3

You can also use:

mask = np.array([[0, 1, 0],
                 [1, 0, 1],
                 [0, 1, 0]]).astype(bool)
a[i-1:i+2, j-1:j+2][mask]

output:

array([53, 93, 94, 86])