Query the value of the four neighbors of an element in a numpy 2D array

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?

>Solution :

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

Leave a Reply