Efficient way of getting the neighbors in 2d numpy array

Question

I want to swap an element in a 3x3 grid with a random neighbour. But of course some elements have 2 neighbours, some 3 and 1 element has 4 neighbours. As I want to choose a random neighbour I have to know how many neighbours each element has.

I'm pretty sure I can do it with a lot if else queries but that's so ugly for this simple problem. But I don't have a different idea how to do it more elegantly.

Answer 1

Embed your 3x3 grid into a 5x5 grid. In the boundary, have some value that will never appear in the 3x3 grid, say -1.

-1 -1 -1 -1 -1 
-1          -1
-1 Your Grid-1
-1          -1
-1 -1 -1 -1 -1

Now, each internal grid cell has 4 neighbors. For these cells, do rejection sampling. That is keep looping a sample from the neighborhood till you draw a sample that is not -1.

Answer 2

Here are two possible solutions.

In both, it is assumed that the position of a cell is a tuple of 0-indexed coordinates (row, col) in a grid.

Solution 1: Clipping

No loop is required and there is 1/n probability to get any neighbour among n possible ones.

import random
import itertools
import numpy as np
# All the possible moves are combination of 2 numbers among 3 (-1, 0 and 1)... 
all_moves = list(itertools.product([-1,0,1], repeat=2))
#... except (0,0) which is keeping the current position.
all_moves.remove((0,0))    

def random_neighbour(current_pos): 
    # Adding them to your current position gives you all virtual neighbours
    choices =  np.array(current_pos) + np.array(all_moves)
    # ... but we must remove the impossible positions
    choices = choices[((choices>=0)&(choices<=2)).all(axis=1)]
    # Now we can choose randomly among the remaining possibilities:
    return random.choice(choices)

Solution 2: Dictionary

This one also insures we get equal probability to get any neighbour to a cell. At the launch of your script, we can initialize a dictionary which provides all neighbours (values) for each cell (keys).

import random
neighbours_dict = {
(p // 3, p % 3): [(p // 3 + x_inc-1, p % 3 + y_inc - 1) 
             for x_inc in range(3) if 1 <= p // 3 + x_inc <= 3
             for y_inc in range(3) if 1 <= p % 3 + y_inc <= 3 and not y_inc == x_inc == 1] 
for p in range(9)}

Then we can simply give it the index of the current position. Please note that positions must be tuples, as lists are not hashable and cannot constitute indices. So we'll cast to make sure we don't get a TypeError.

random.choice(neighbours_dict[tuple(current_pos)])

Better ways to initialise the dictionary

Of course, we can also manually initialize the dictionary to fit our specific needs, like so:

neighbours_dict = {(0, 0): [(0, 1), (1, 0), (1, 1)],
 (1, 0): [(0, 0), (0, 1), (1, 1), (2, 0), (2, 1)],
 (2, 0): [(1, 0), (1, 1), (2, 1)],
 (0, 1): [(0, 0), (0, 2), (1, 0), (1, 1), (1, 2)],
 (1, 1): [(0, 0), (0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1), (2, 2)],
 (2, 1): [(1, 0), (1, 1), (1, 2), (2, 0), (2, 2)],
 (0, 2): [(0, 1), (1, 1), (1, 2)],
 (1, 2): [(0, 1), (0, 2), (1, 1), (2, 1), (2, 2)],
 (2, 2): [(1, 1), (1, 2), (2, 1)]}

Or we can generalize the dictionary comprehension to make it compatible to any grid shape and size:

grid_width, grid_height = 3, 3
neighbours_dict = {
    (p // grid_width, p % grid_width):
        [(p // grid_width + y_inc - 1, p % grid_width + x_inc - 1) 
            for y_inc in range(3) if 1 <= p // grid_width + y_inc <= grid_height
            for x_inc in range(3) if 1 <= p % grid_width + x_inc <= grid_width 
                    and not y_inc == x_inc == 1] 
    for p in range(grid_width * grid_height) }

This latter is, for me, the most elegant solution of those.