Найти сопоставление индексов между двумя массивами numpy

Question

Есть ли хороший способ в numpy, чтобы получить индексные данные о том, где каждый элемент в array1 находится в array2?

Пример:

array1 = np.array([1, 3, 4]) 
array2 = np.arange(-2, 5, 1, dtype=np.int) 

np.where(array1[0] == array2) 
# (array([3]),) 
np.where(array1[1] == array2) 
# (array([5]),) 
np.where(array1[2] == array2) 
# (array([6]),) 

Я хотел бы сделать

np.where(array1 == array2) 
# (array([3 5 6]),) 

ли что-то подобное возможно? Мы гарантируем, что все записи в array1 могут быть найдены в array2.

Answer 1

Подход № 1: Используйте np.in1d там, чтобы получить маску местах, где происходят матчи, а затем np.where, чтобы получить эти позиции индекса -

np.where(np.in1d(array2, array1)) 

Подход № 2: С np.searchsorted -

np.searchsorted(array2, array1) 

Обратите внимание, что если array2 не отсортировано, нам необходимо использовать дополнительный необязательный аргумент sorter с ним.

Пример запуска -

In [14]: array1 
Out[14]: array([1, 3, 4]) 

In [15]: array2 
Out[15]: array([-2, -1, 0, 1, 2, 3, 4]) 

In [16]: np.where(np.in1d(array2, array1)) 
Out[16]: (array([3, 5, 6]),) 

In [17]: np.searchsorted(array2, array1) 
Out[17]: array([3, 5, 6]) 

Продолжительность испытания -

In [62]: array1 = np.random.choice(10000,1000,replace=0) 

In [63]: array2 = np.sort(np.random.choice(100000,10000,replace=0)) 

In [64]: %timeit np.where(np.in1d(array2, array1)) 
1000 loops, best of 3: 483 µs per loop 

In [65]: %timeit np.searchsorted(array2, array1) 
10000 loops, best of 3: 40 µs per loop