如何从一维数组中获取二维数组索引?

时间:2019-01-25 20:14:54

标签: python numpy

我正在寻找一种基于1d数组中的值返回2d数组索引的有效方法。我目前有一个嵌套的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的索引。我想要的输出将是这个2d数组:

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网格点。

2 个答案:

答案 0 :(得分:1)

使用np.searchsorted

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

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

searchsorted使用复杂的data2d执行二进制搜索。然后重新调整结果。


另一个选择是建立索引并在固定时间内查询它。您可以使用熊猫的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]])

答案 1 :(得分:0)

这也应该很快

import numpy as np
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])
idxdict = dict(zip(data1d,range(len(data1d))))
locs = data2d
for i in range(len(locs)):
    for j in range(len(locs[i])):
        locs[i][j] = idxdict[locs[i][j]]