我正在使用Python中的n维数组,我希望根据其坐标找到给定单元格的“邻居”(邻近单元格)。问题是我事先并不知道维度的数量。
我尝试按numpy.roll
的建议使用this answer,但似乎不清楚如何将此方法应用于多个维度。
请指出正确的方向。
答案 0 :(得分:2)
我将假设你有一个(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]]
这将推广到任何维度。
如果您只想索引&#34;邻居&#34; p
并且您并不关心排除p
本身,更简单,更快捷的选择是使用slice
个对象的元组:
idx = tuple(slice(pp - 1, pp + 2) for pp in p)
print(x[idx])
# [[1 1]
# [1 2]
# [1 1]]