在numpy ndarray中查找邻居

Question

我遇到了一个问题，我有一个numpy的ndarray数组。这个数组可以是二维的，也可以是三维的，这取决于输入的数据。我想要做的是获取每个样本在每个轴上的“前一个”邻居。所以在二维情况下，假设有这样一个数组：

1 2 3 4
5 6 7 8
4 3 2 1

我只需要那些实际上有有效“前一个”邻居的值的索引（也就是说，不包括顶部和左边的边界项）。所以这应该返回我一些像这样的结果：

[[0, 1], [1, 0]], [[1, 1], [2, 0]], [[2, 1], [3, 0]] # corresponding to 6, 7, 8 entries
[[0, 2], [1, 1]], [[1, 2], [2, 1]], [[2, 2], [3, 1]] # corresponding to 3, 2, 1 entries

基本上只考虑子数组。

6 7 8
3 2 1

另外，如果是三维数组，那么还有一个轴，但规则类似，也就是只获取有效条目的子数组。我想用这些索引来操作原始数组中的条目。不过，我不知道怎么高效地生成这些索引，而不需要在ndarray的轴上进行循环（而且这个轴的数量也不是固定的）。

Answer 1

考虑一下你的数据是这样的：

import numpy as np
data = np.array([[1, 2, 3, 4],
                 [5, 6, 7, 8],
                 [4, 3, 2, 1]])

如果我理解正确的话，使用切片（slicing）可以得到你需要的结果，也就是：

>>> data[1:, 1:]
array([[6, 7, 8],
       [3, 2, 1]])

因此，如果有效的前一个邻居意味着这些项目的上面和左边都有邻居，你只需对数组的每个维度应用相同的切片 1:。在三维数据中，可以写成 data[1:, 1:, 1:] 等等。

Answer 2

好的，我搞明白了。在我的情况下，我可以简单地使用numpy的索引。所以可以这样写：

import numpy as np
i = np.indices(x-2 for x in data.shape)

然后我可以这样使用它：

data[i[0], i[1], i[2]...]