如何在 ndarray 中找到单元格的邻居？答案

【问题标题】：How to Find the Neighbors of a Cell in an ndarray?如何在 ndarray 中找到单元格的邻居？
【发布时间】：2016-04-26 14:35:05
【问题描述】：

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

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

请指点我正确的方向。

【问题讨论】：

myarray.shape 保存有关数组维度的信息
@yurib 我知道这一点。我的意思是我不能将实现建立在一定数量的维度上。

标签： python arrays numpy multidimensional-array scipy

【解决方案1】：

我将假设您有一个 (ndims,) 索引向量指定某个点 p，并且您想要一个与数组中每个相邻元素的位置相对应的索引 (m, ndims) 数组（包括对角线相邻元素）。

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

您可能希望排除点 p 本身（即 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

这是一个易于可视化的 2D 示例：

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]]

【讨论】：