python之如何在 ndarray 中查找单元格的邻居

我在 Python 中使用 n 维数组，我想根据给定单元格的坐标找到它的“邻居”(相邻单元格)。问题是我事先不知道维数。

我尝试使用 numpy.roll正如 this answer 所建议的那样, 但似乎不清楚如何将此方法应用于多个维度。

请指出正确的方向。

请您参考如下方法:

我假设你有一个 (ndims,) 索引向量指定了一些点 p，你想要一个 (m, ndims ) 索引数组对应于数组中每个相邻元素的位置(包括对角线相邻元素)。

从索引向量 p 开始，您希望每个元素偏移 -1、0 和 +1 的所有可能组合。这可以通过使用 np.indices 来完成生成一个 (m, ndims) 偏移量数组，然后将这些偏移量添加到 p。

您可能想要排除点 p 本身(即 where offset == np.array([0, 0, ..., 0])，并且您可能还需要排除越界索引。

import numpy as np 
 
def get_neighbours(p, exclude_p=True, shape=None): 
 
    ndim = len(p) 
 
    # generate an (m, ndims) array containing all strings over the alphabet {0, 1, 2}: 
    offset_idx = np.indices((3,) * ndim).reshape(ndim, -1).T 
 
    # use these to index into np.array([-1, 0, 1]) to get offsets 
    offsets = np.r_[-1, 0, 1].take(offset_idx) 
 
    # optional: exclude offsets of 0, 0, ..., 0 (i.e. p itself) 
    if exclude_p: 
        offsets = offsets[np.any(offsets, 1)] 
 
    neighbours = p + offsets    # apply offsets to p 
 
    # optional: exclude out-of-bounds indices 
    if shape is not None: 
        valid = np.all((neighbours < np.array(shape)) & (neighbours >= 0), axis=1) 
        neighbours = neighbours[valid] 
 
    return neighbours 

这是一个易于可视化的二维示例:

p = np.r_[4, 5] 
shape = (6, 6) 
 
neighbours = get_neighbours(p, shape=shape) 
 
x = np.zeros(shape, int) 
x[tuple(neighbours.T)] = 1 
x[tuple(p)] = 2 
 
print(x) 
# [[0 0 0 0 0 0] 
#  [0 0 0 0 0 0] 
#  [0 0 0 0 0 0] 
#  [0 0 0 0 1 1] 
#  [0 0 0 0 1 2] 
#  [0 0 0 0 1 1]] 

这将推广到任何维度。

如果您只想索引 p 的“邻域”并且您不关心排除 p 本身，那么一个更简单和更快的选择是是使用 slice 对象的元组:

idx = tuple(slice(pp - 1, pp + 2) for pp in p) 
print(x[idx]) 
# [[1 1] 
#  [1 2] 
#  [1 1]]