Python，规则网格上的邻居

Question

假设我有一组二维坐标，它们代表二维规则网格的单元中心。我想找出，对于网格中的每个单元，每个方向上的两个最近的邻居。在

如果将每个单元格和索引分配给如下定义，则问题非常简单：

idx_cell=idx+N*ide

其中N是网格中的单元格总数，idx=x/dx，idy=y/dx，x和y是单元格的x坐标和y坐标，dx是单元格的大小。在

例如，idx_cell=5的单元的相邻单元是idx_cell等于4,6（对于x轴）和5+N，5-N（对于y轴）的单元。在

我的问题是，对于大型（N>；1e6）数据集，我的算法实现非常慢。在

例如，为了得到x轴的邻域

[x[(idx_cell==idx_cell[i]-1)|(idx_cell==idx_cell[i]+1)] for i in cells]

你认为有没有最快的方法来实现这个算法？在

Answer 1

您基本上是在重新设计多维数组的索引方案。编写代码相对容易，但是您可以在这里使用两个函数^{}和{a2}。在

如果网格包含M行和N列，要获得单个项目的idx和{}，可以执行以下操作：

>>> M, N = 12, 10
>>> np.unravel_index(4, dims=(M, N))
(0, 4)

如果您提供一个索引数组而不是单个索引，则此方法也有效：

^{pr2}$

因此，如果cells有多个单元格的索引，则您希望找到其邻居：

^{3}$

你可以让他们的邻居：

>>> idy, idx = np.unravel_index(cells, dims=(M, N))
>>> neigh_idx = np.vstack((idx-1, idx+1, idx, idx))
>>> neigh_idy = np.vstack((idy, idy, idy-1, idy+1))
>>> np.ravel_multi_index((neigh_idy, neigh_idx), dims=(M,N))
array([[14, 27, 31, 43, 86],
       [16, 29, 33, 45, 88],
       [ 5, 18, 22, 34, 77],
       [25, 38, 42, 54, 97]], dtype=int64)

或者，如果你喜欢这样：

>>> np.ravel_multi_index((neigh_idy, neigh_idx), dims=(M,N)).T
array([[14, 16,  5, 25],
       [27, 29, 18, 38],
       [31, 33, 22, 42],
       [43, 45, 34, 54],
       [86, 88, 77, 97]], dtype=int64)

这样做的好处是ravel_multi_index有一个mode关键字参数，可以用来处理晶格边缘上的项，请参阅文档。在