Looking For A Readable, Elegant Way To Select Direct Neighbors In 2d List In Python

I have 2D list that I want to obtain direct neighbors from (up,down,left,right) and I was wondering what's the most pythonic way of doing this. I've looked at Determining neighbour

Solution 1:

If you want to use numpy, you can use an indexing array, which contains the indices of the neighbors relative to the index you want, and then add this to your desired index. I personally think this is elegant, but YMMV

Here's an example:

import numpy as np

# A 5 * 5 grid
grid = np.arange(25).reshape(5, 5)

# A schematic representation of the grid# 0,  1,  2,  3,  4# 5,  6,  7,  8,  9# 10, 11, 12, 13, 14# 15, 16, 17, 18, 19# 20, 21, 22, 23, 24# We define how our neighbors relate to our index.
mask = np.array([[0, 1], [1, 0], [0, -1], [-1, 0]])

# Let's say we want the neighbors of [2, 2], which are 17, 11, 7, and 13# Index marked with X, neighbors marked with O# 0,  1,  2,  3,  4# 5,  6,  O   8,  9# 10, O   X   O  14# 15, 16, O   18, 19# 20, 21, 22, 23, 24

desired_index = np.array([2, 2])

# We add the neighbor indices to the mask
neighbor_indices = desired_index + mask
# [[2, 3], [3, 2], [2, 1], [1, 2]]# Index the array using the indices.
neighbors = grid[neighbor_indices[:, 0], neighbor_indices[:, 1]]

Note that this example does not take care of out of bounds issues. Specifically, it will error when given indices higher than the number of columns or rows, and will wrap around for indices < 0.

desired_index = np.array([0, 0])
neighbor_indices = desired_index + mask
neighbors = grid[neighbor_indices[:, 0], neighbor_indices[:, 1]]
# Wrong

desired_index = np.array([4, 4])
neighbor_indices = desired_index + mask
neighbors = grid[neighbor_indices[:, 0], neighbor_indices[:, 1]]
# Error