Python，规则网格上的邻居

假设我有一组二维坐标，表示二维规则网格单元的中心。我想为网格中的每个单元找到每个方向上最近的两个邻居

如果将每个单元格和索引指定为以下定义，则问题非常简单：

idx_单元=idx+N*idy

其中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_单元==idx_单元[i]-1）|（idx_单元==idx_单元[i]+1）]用于单元中的i]

---

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

你基本上是在重新设计多维数组的索引方案。编写代码相对容易，但您可以使用这两个函数，并在这里发挥您的优势

如果网格由

M

行和

N

列组成，要获取单个项目的

idx

和

idy

，可以执行以下操作：

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

如果您提供了一个索引数组而不是单个索引，则也可以这样做：

>>> np.unravel_index([15, 28, 32, 97], dims=(M, N))
(array([1, 2, 3, 9], dtype=int64), array([5, 8, 2, 7], dtype=int64))

因此，如果

单元

具有要查找相邻单元的多个单元的索引：

>>> cells = np.array([15, 28, 32, 44, 87])

您可以通过以下方式获取他们的邻居：

>>> 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\u multi\u index

有一个

mode

关键字参数，可以用来处理晶格边缘上的项目，请参见文档。

我认为可以通过使用NumPy fancy索引来加快速度。你可以发布更多的代码来创建idx，idy，cells，idx_cells，x。你也可以使用scipy的cKDTree。谢谢，我已经试过使用KDTree，但是它对于我的数据集来说相当慢，而且对于网格来说也不方便。不幸的是，我无法发布更多的代码，因为我的数据集太大了。