如何在ndarray中查找单元格的邻居? [英] How to Find the Neighbors of a Cell in an ndarray?

问题描述

我正在使用Python中的n维数组，并且我想根据其坐标找到给定单元的邻居"(相邻单元).问题是我事先不知道尺寸数.

I'm working with n-dimensional arrays in Python, and I want to find the "neighbors" (adjacent cells) of a given cell based on its coordinates. The issue is that I don't know the number of dimensions beforehand.

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

I attempted to use numpy.roll as suggested by this answer, but it seems unclear how to apply this method to multiple dimensions.

请指向正确的方向.

推荐答案

我将假设您有一个索引(ndims,)向量，该向量指定了某个点p，并且您需要一个(m, ndims)索引数组对应于数组中每个相邻元素(包括对角相邻元素)的位置.

I'm going to assume that you have an (ndims,) vector of indices specifying some point p, and you want an (m, ndims) array of indices corresponding to the locations of every adjacent element in the array (including diagonally adjacent elements).

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

Starting out with your indexing vector p, you want to offset each element by every possible combination of -1, 0 and +1. This can be done by using np.indices to generate an (m, ndims) array of offsets, then adding these offsets to p.

您可能要排除点p本身(即offset == np.array([0, 0, ..., 0])所在的位置)，并且可能还需要排除边界索引.

You might want to exclude point p itself (i.e. where offset == np.array([0, 0, ..., 0]), and you may also need to exclude out-of-bounds indices.

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示例:

Here's a 2D example that's easy to visualize:

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对象的元组:

If you just want to be able to index the "neighbourhood" of p and you don't care about excluding p itself, a much simpler and faster option would be to use a tuple of slice objects:

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

这篇关于如何在ndarray中查找单元格的邻居?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！