python - 如何从一维数组获取二维数组的索引？

Question

我正在寻找一种有效的方法来根据一维数组中的值返回二维数组的索引。我目前设置了一个嵌套的 for 循环，速度非常慢。

这里是一些示例数据和我想要得到的:

data2d = np.array( [  [1,2] , [1,3] ,[3,4], [1,2] , [7,9] ])

data1d = np.array([1,2,3,4,5,6,7,8,9])

我想返回 data2d 等于 data1d 的索引。我想要的输出是这个二维数组:

locs = np.array([[0, 1], [0, 2], [2, 3], [0, 1], [6, 8]])

我唯一想到的是嵌套 for 循环:

locs = np.full((np.shape(data2d)), np.nan)

for i in range(0, 5):
    for j in range(0, 2):
        loc_val = np.where(data1d == data2d[i, j])
        loc_val = loc_val[0]
        locs[i, j] = loc_val

这对于一小部分数据来说没问题，但我有 87,600 个 2d 网格，每个网格点都是 428x614。

Answer 1

使用np.searchsorted :

np.searchsorted(data1d, data2d.ravel()).reshape(data2d.shape)

array([[0, 1],
       [0, 2],
       [2, 3],
       [0, 1],
       [6, 8]])

searchsorted 使用 ravelled data2d 执行二分搜索。然后重新调整结果。

<小时/>

另一种选择是构建索引并在恒定时间内对其进行查询。您可以使用 pandas 的 Index API 来完成此操作。

import pandas as pd

idx = pd.Index([1,2,3,4,5,6,7,8,9])
idx
#  Int64Index([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype='int64')

idx.get_indexer(data2d.ravel()).reshape(data2d.shape)

array([[0, 1],
       [0, 2],
       [2, 3],
       [0, 1],
       [6, 8]])