Question

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

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

idx_cell = idx + N * idy

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

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

我遇到的问题是，对于大型（N＆gt; 1e6）数据集，我的算法实现速度很慢。

例如，为了获得x轴的邻居，我做了

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

您认为有最快的方法来实现此算法吗？

Answer 1

您基本上正在重新发明多维数组的索引方案。代码相对容易，但您可以使用这两个函数unravel_index和ravel_multi_index。

如果您的网格属于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具有几个单元格的索引，则要查找邻居：

>>> 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_multi_index有一个mode关键字参数，可用于处理晶格边缘的项目，请参阅文档。

Python，常规网格上的邻居

1 个答案: