4

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

如果将每个单元格和索引分配给定义如下：

idx_cell = 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_cell==idx_cell[i]-1)|(idx_cell==idx_cell[i]+1)] for i in cells]

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

4

1 回答 1

4

您基本上是在重新发明多维数组的索引方案。编码相对容易，但您可以在这里使用这两个功能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关键字参数，您可以使用它来处理格子边缘的项目，请参阅文档。

于 2013-03-28T17:15:04.757 回答