如何切割2D numpy数组以获得其直接邻居？

Question

I want to loop through my 2D numpy array and check all its direct neighbors. If I make a numpy array like so:

tilemap = np.arange(16).reshape(4,4)

It will look something like this:

 [[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

The loop that I created to help me find the neighbors for each spot in the array looks like this:

 import numpy as np

 mapwidth = 4
 mapheight = 4

 tilemap = np.arange(mapwidth * mapheight).reshape(mapwidth, mapheight)

 row = 0
 for i in tilemap:
     count = 0
     for j in i:
         column = j % mapwidth
         check = tilemap[row-1:row+2, column-1:column+2]
         print(check)
         count += 1
         if count % mapheight == 0:
             row += 1

However, when I do this, I will not find any neighbors for the spot in the array with the values 0, 1, 2, 3, 4, 8, and 12. I understand why this is the case. for example, if I take the value of 8. it has the indices of [2,0]. the row-1 will result in -1, which is the same as index 3 in this example. the row+2 is 2. And slicing 2:3 will result in nothing because there is nothing between 2 and 3.

Anyway, the result I'm looking for is something like this (for value 8):

[[4  5]
[ 8  9]
[12 13]]

I know I can achieve this by piling some if-statements but I wonder if there is a more elegant way of dealing with this.

Thanks for your time.

(For those who are curious to know this): The neighbors for eg value 11 actually return like I want them to, without any errors. It returns this:

[[6  7]
[10 11]
[14 15]]

EDIT:

I should also mention I tried this:

check = np.take(tilemap, tilemap[row-1:row+2, column-1:column+2], mode = 'clip')

But this did not work.

Answer 1

You can simplify the way your loop is written and not assume as much about the contents of the array, making your code more flexible. Numpy has an nditer class that can be used to iterate over an array. You can also use it to get the multi-dimensional index of each element. Iteration can be further simplified using the ndenumerate class, similar to Pythons builtin enumerate . If you do not need to get the elements back, just the index, you can use ndindex . Here is an example using ndindex :

for r, c in ndindex(tilemap.shape):
    check = tilemap[max(r-1, 0):min(r+1, mapheight), max(c-1, 0):min(c+1, mapwidth)]
    print(check)