Python find list of surrounding neighbours of a node in 2D array

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, i.e. 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)]