Python在2D数组中查找节点周围邻居的列表

Question

I've been working on a code (Py 2.7) that generates an array of elements with each node assigned some random numbers. Now, I wish to make a list of the surrounding elements, and find the index of the max value. The array size is variable (I considered col = array column size). I have assigned numbers to each node (I called it 's' in the below) so that I can find the 2D index of the array element. Here is what I wrote

rn = s/col; cn = s%col;
b = [rr[rn,cn+1],rr[rn-1,cn+1],rr[rn-1,cn],rr[rn-1,cn-1],rr[rn,cn-1],rr[rn+1,cn-1],rr[rn+1,cn],rr[rn+1,cn+1]]
ma = max(b)
a = [i for i,j in enumerate(b) if j == ma]

Is there any short method to find the neighbours without the need to number each array element ? (like I did using s).

Answer 1

You can use numpy for this. First, let's create a random 5x5 matrix M for testing...

>>> M = np.random.random((5, 5))
>>> M
array([[ 0.79463434,  0.60469124,  0.85488643,  0.69161242,  0.25254776],
       [ 0.07024954,  0.84918038,  0.01713536,  0.42620873,  0.97347887],
       [ 0.3374191 ,  0.99535699,  0.79378892,  0.0504229 ,  0.05136649],
       [ 0.73609556,  0.94250215,  0.67322277,  0.49043047,  0.60657825],
       [ 0.71153444,  0.43242926,  0.29726895,  0.2173065 ,  0.38457722]])

Now we take a slice from this matrix, N , holding the neighbors of some central element (x, y)

>>> x, y = 2, 2
>>> N = M[x-1:x+2, y-1:y+2]
>>> N
array([[ 0.84918038,  0.01713536,  0.42620873],
       [ 0.99535699,  0.79378892,  0.0504229 ],
       [ 0.94250215,  0.67322277,  0.49043047]])

We can now get a new matrix showing which of the elements of the orginal matrix M is equal to the max from N

>>> M == N.max()
array([[False, False, False, False, False],
       [False, False, False, False, False],
       [False,  True, False, False, False],
       [False, False, False, False, False],
       [False, False, False, False, False]], dtype=bool)

Now we can use numpy.where to get the index of the element(s) that are True in this matrix. zip those to get a list of tuples.

>>> zip(*np.where(M == N.max()))
[(2, 1)]

Note that those are the positions in the original matrix M , ie they could contain tuples that are not in N . Alternatively, you can get just the maximum values in N , but then you'll have to add x-1 and y-1 as offset afterwards.

>>> zip(*np.where(N == N.max()))
[(1, 0)]