Получение индексов нескольких элементов в массиве NumPy сразу

Есть ли способ получить индексы нескольких элементов в массиве NumPy сразу?

например.

import numpy as np
a = np.array([1, 2, 4])
b = np.array([1, 2, 3, 10, 4])

Я хотел бы найти индекс каждого элемента a в b, а именно: [0,1,4].

Я нахожу решение, я использую бит многословный:

import numpy as np

a = np.array([1, 2, 4])
b = np.array([1, 2, 3, 10, 4])

c = np.zeros_like(a)
for i, aa in np.ndenumerate(a):
    c[i] = np.where(b==aa)[0]

print('c: {0}'.format(c))

Вывод:

c: [0 1 4]

Ответы

Ответ 1

Вы можете использовать in1d и nonzero (или where, если на то пошло):

>>> np.in1d(b, a).nonzero()[0]
array([0, 1, 4])

Это отлично подходит для ваших массивов примеров, но в целом массив возвращаемых индексов не выполняет порядок значений в a. Это может быть проблемой в зависимости от того, что вы хотите сделать дальше.

В этом случае гораздо лучший ответ - тот, который @Jaime дает здесь, используя searchsorted:

>>> sorter = np.argsort(b)
>>> sorter[np.searchsorted(b, a, sorter=sorter)]
array([0, 1, 4])

Это возвращает индексы для значений, как они отображаются в a. Например:

a = np.array([1, 2, 4])
b = np.array([4, 2, 3, 1])

>>> sorter = np.argsort(b)
>>> sorter[np.searchsorted(b, a, sorter=sorter)]
array([3, 1, 0]) # the other method would return [0, 1, 3]

Ответ 2

Это простой однострочный файл с использованием пакета numpy-indexed (отказ от ответственности: я являюсь его автором):

import numpy_indexed as npi
idx = npi.indices(b, a)

Реализация полностью векторизована, и она дает вам контроль над обработкой отсутствующих значений. Более того, он работает и для nd-массивов (например, нахождение индексов строк a в b).