Looking for a readable, elegant way to select direct neighbors in 2D list in Python

Question

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 neighbours of cell two dimensional list but their one soltution to finding direct neighbors isn't doing it for me: (let x,y be any two indices in the 2D list)

neighbors = [(x+a[0], y+a[1]) for a in 
                    [(-1,0), (1,0), (0,-1), (0,1)]

If anything, I would have done it like this:

neighbors = [(x+a,y+b) for a,b in 
                    [(-1,0), (1,0), (0,-1), (0,1)]

or like this:

neighbors = [(a,b) for a,b in 
                    [(x-1,y), (x+1,y), (x,y-1), (x,y+1)]

but the latter feels a bit hardcoded. Thoughts?

EDIT: To formalize my question: what is a readable, elegant way to get direct neighbors from a 2D list in Python?

I don't understand; what is your question? You want less code or more code? — Ralf, Jul 03 '19 at 12:24
I want something readable; I don't mind writing more code if it means others can understand what I'm trying to do. — Jonathan Esteban, Jul 03 '19 at 12:35
If you play with big 2D or image maybe look into convolution method, to improve speed — themadmax, Jul 03 '19 at 12:38
`x` and `y` are undefined here: `neighbors = [(x,y) for x,y in [(x-1,y), (x+1,y), (x,y-1), (x,y+1)]]` — Reblochon Masque, Jul 03 '19 at 12:39
Assume x,y to be defined as any two valid indices of the 2D list. I updated the code to make it less confusing. — Jonathan Esteban, Jul 03 '19 at 12:47

score 2 · Accepted Answer · answered Jul 03 '19 at 14:06

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

That's amazing. Very elegant indeed. Although I'm only working with small, plain 2D lists so using numpy is a bit of an overkill. My question is really more about expressing my intents clearly. Maybe I'm just overthinking it. — Jonathan Esteban, Jul 03 '19 at 14:49
If you think the array is overkill you could still take a similar approach using lists. `[grid[desired_x + x, desired_y + y] for x, y in [[-1, 0], [0, -1], [1, 0], [0, 1]]]` It looks a bit clunky, but I think you get the idea. In any case, I thought the original idea was clear. — amdex, Jul 03 '19 at 14:55

Looking for a readable, elegant way to select direct neighbors in 2D list in Python

1 Answers1

Linked