在numpy数组中获取一个元素相邻的元素

Question

我有一个类似这样的数组：

[['A0' 'B0' 'C0']
 ['A1' 'B1' 'C1']
 ['A2' 'B2' 'C2']]

我希望获取B1的邻居节点，它们是B0，C1，B2，A1，以及它们的索引。以下是我的实现：

import numpy as np


arr = np.array([
    ['A0','B0','C0'],
    ['A1','B1','C1'],
    ['A2','B2','C2'],
])


def get_neighbor_indices(x,y):
    neighbors = []
    try:
        top = arr[y - 1, x]
        neighbors.append((top, (y - 1, x)))
    except IndexError:
        pass
    try:
        bottom = arr[y + 1, x]
        neighbors.append((bottom, (y + 1, x)))
    except IndexError:
        pass
    try:
        left = arr[y, x - 1]
        neighbors.append((left, (y, x - 1)))
    except IndexError:
        pass
    try:
        right = arr[y, x + 1]
        neighbors.append((right, (y, x + 1)))
    except IndexError:
        pass
    return neighbors

这将返回一个元组列表 (value, (y, x))。

是否有更好的方法来做到这一点，而不依赖于 try/except 呢？

- conquistador

---

实际上，arr[-1, 0]会给你A2。你想忽略这种情况吗？ - Yaroslav Kornachevskyi

Answer 1

你可以直接使用numpy进行操作，不需要任何异常处理，因为您知道数组的大小。即刻邻居的索引x, y由以下给出：

inds = np.array([[x, y]]) + np.array([[1, 0], [-1, 0], [0, 1], [0, -1]])

您可以轻松地创建一个指示哪些索引是有效的掩码：

valid = (inds[:, 0] >= 0) & (inds[:, 0] < arr.shape[0]) & \
        (inds[:, 1] >= 0) & (inds[:, 1] < arr.shape[1])

现在提取您想要的值：

inds = inds[valid, :]
vals = arr[inds[:, 0], inds[:, 1]]

最简单的返回值是 inds, vals，但如果您坚持保留原始格式，可以将其转换为：

[v, tuple(i) for v, i in zip(vals, inds)]

附录

您可以轻松修改此代码以在任意维度上运行：

def neighbors(arr, *pos):
    pos = np.array(pos).reshape(1, -1)
    offset = np.zeros((2 * pos.size, pos.size), dtype=np.int)
    offset[np.arange(0, offset.shape[0], 2), np.arange(offset.shape[1])] = 1
    offset[np.arange(1, offset.shape[0], 2), np.arange(offset.shape[1])] = -1
    inds = pos + offset
    valid = np.all(inds >= 0, axis=1) & np.all(inds < arr.shape, axis=1)
    inds = inds[valid, :]
    vals = arr[tuple(inds.T)]
    return vals, inds

给定一个N维数组arr和N个元素的pos，您可以通过依次将每个维度设置为1或-1来创建偏移量。通过在inds和arr.shape之间进行广播，并在每个大小为N的行上调用np.all，而不是手动为每个维度执行此操作，可以大大简化掩码valid的计算。最后，将inds转换为实际的花式索引，方法是将每列分配到单独的维度，即tuple(inds.T)。转置是必要的，因为数组沿着行（dim 0）迭代。

- Mad Physicist

---

我将掩码更改为[[-1, 0], [0, 1], [1, 0], [0, -1]]以对应于上、右、下、左。您的掩码给了我左上角、右上角、右下角、左下角。 - conquistador

@conquistador，做得好。我写那个的时候想的是什么我自己也不知道。现在已经修正了。 - Mad Physicist

Answer 2

你可以使用这个：

def get_neighbours(inds):
    places = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    return [(arr[x, y], (y, x)) for x, y in [(inds[0] + p[0], inds[1] + p[1]) for p in places] if x >= 0 and y >= 0]

get_neighbours(1, 1)
# OUTPUT [('B0', (1, 0)), ('B2', (1, 2)), ('A1', (0, 1)), ('C1', (2, 1))]

get_neighbours(0, 0)
# OUTPUT [('A1', (0, 1)), ('B0', (1, 0))]

- Aryerez

Answer 3

这个怎么样？

def get_neighbor_indices(x,y):
    return ( [(arr[y-1,x], (y-1, x))] if y>0 else [] ) + \
           ( [(arr[y+1,x], (y+1, x))] if y<arr.shape[0]-1 else [] ) + \
           ( [(arr[y,x-1], (y, x-1))] if x>0 else [] ) + \
           ( [(arr[y,x+1], (y, x+1))] if x<arr.shape[1]-1 else [] )

- Andrew Kay