Warm tip: This article is reproduced from stackoverflow.com, please click
python

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

发布于 2020-03-27 10:21:06

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?

Questioner
Jonathan Esteban
Viewed
18
amdex 2019-07-03 22: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