CODEWITHSUNDEEP
pythonnumpy
skywave1980
skywave1980

Reputation: 109

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?

Upvotes: 1

Views: 759

Answers (3)

Kyriakos Psarakis
Kyriakos Psarakis

Reputation: 84

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 = np.array(matrix.shape) - 1
    nbr = []
    if x > x_len or y > y_len:
        return 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

Upvotes: 1

PlainRavioli
PlainRavioli

Reputation: 1211

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

Upvotes: 1

mozway
mozway

Reputation: 260420

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

Upvotes: 1

Related Questions