Numpy implementation can take advantage numpy.argwhere to retrieve value indeces, create the grid of indices with numpy.ix_, and finally apply numpy.narray.ravel method to flat the array::
import numpy as npn = 5grid = np.arange(n**2).reshape(n,n)[::-1]def celdas_vecinas_np(grid, v, n): x, y = np.argwhere(grid == v)[0] idx = np.arange(x-1, x+2) %n idy = np.arange(y-1, y+2) %n return grid[np.ix_(idx, idy)].ravel()celdas_vecinas_np(grid, 24, n)array([ 3, 4, 0, 23, 24, 20, 18, 19, 15])On the other hand, for a Numba implementation we cannot use numpy.argwhere, but we can use numpy.where to get indices. Once we do that, it is only a matter of looping in the right ranges, namely:
from numba import njit@njitdef celdas_vecinas_numba(grid, v, n): x, y = np.where(grid == v) x, y = x[0], y[0] result = [] for ix in range(x-1, x+2): for iy in range(y-1, y+2): result.append(grid[ix%n, iy%n]) return resultceldas_vecinas_numba(grid, 24, n)[3, 4, 0, 23, 24, 20, 18, 19, 15]Perfomance Comparison with such small grid numba already runs ~20x faster in my local machine:
%timeit celdas_vecinas_np(grid, 24, 5)38 µs ± 1.62 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)%timeit celdas_vecinas_numba(grid, 24, n)1.81 µs ± 93.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)