如何在ndarray中查找單元格的鄰居？

Question

我正在Python中使用n維數組，並且我想根據其座標找到給定單元格的「鄰居」（相鄰單元格）。問題是我不知道預先的維數。

我試圖按照this answer的建議使用numpy.roll，但似乎不清楚如何將此方法應用於多個維度。

請指點我正確的方向。

Answer 1

我要去假定有索引指定一些點p的(ndims,)矢量，並且想要的(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 combinations of 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]]