Cilenco
Cilenco

Reputation: 7117

Numpy get neighbors always as 3x3 matrix

Let's say I have a 2d numpy array of the size (5,5). I can get the neighbors of the index (i,j) with the following statement:

a = range(25)
a = np.reshape(a, (5,5))

n = a[i-1:i+2, j-1:j+2]

That works great for 0 < i,j < 4. My problem is that I always want to get a 3x3 array but if one of the indices is 0 or 4 I do not get it (in case i=0 the range is (-1, 2) = (4, 2) and we get an empty range)

Do you have any ideas how I can always get a 3x3 matrix and fill the "failed indices" with zeros?

Upvotes: 3

Views: 1136

Answers (1)

YXD
YXD

Reputation: 32511

Use np.pad to extend your array first and index into the result instead. You'll have to shift your indices accordingly.

>>> b = np.pad(a, pad_width=1, mode='constant')
>>> b
array([[ 0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  1,  2,  3,  4,  0],
       [ 0,  5,  6,  7,  8,  9,  0],
       [ 0, 10, 11, 12, 13, 14,  0],
       [ 0, 15, 16, 17, 18, 19,  0],
       [ 0, 20, 21, 22, 23, 24,  0],
       [ 0,  0,  0,  0,  0,  0,  0]])
>>>

Upvotes: 4

Related Questions