时间戳

首页 > 解决方案 > 如何从一维数组中获取二维索引数组？

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

问题描述

我正在寻找一种有效的方法来根据一维数组中的值返回二维数组的索引。我目前有一个嵌套的 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 网格点。

标签： pythonnumpy

解决方案

使用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 的IndexAPI 来做到这一点。

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]])

推荐阅读