Как получить 2d массив индексов из 1d массива?

Question

Я ищу эффективный способ вернуть индексы для двумерного массива на основе значений в одномерном массиве. В настоящее время у меня мучительно медленная установка вложенного цикла.

Вот некоторые примеры данных и что я хочу получить:

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

Единственное, что я придумал, это вложенный цикл:

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 двумерных сеток, каждая из которых имеет 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 выполняет бинарный поиск с равниной data2d. Результат затем изменяется.

---

Другой вариант - создать индекс и запросить его за постоянное время. Вы можете сделать это с помощью Index API pandas.

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

Answer 2

Это должно быть быстро также

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