Найти несколько значений в массиве Numpy

Question

Я ищу функцию numpy, чтобы найти индексы, при которых определенные значения находятся внутри вектора (xs).Значения приведены в другом массиве (ys).Возвращаемые индексы должны следовать порядку ys.

В коде я хочу заменить приведенное ниже понимание списка на функцию numpy.

>> import numpy as np
>> xs = np.asarray([45, 67, 32, 52, 94, 64, 21])
>> ys = np.asarray([67, 94])
>> ndx = np.asarray([np.nonzero(xs == y)[0][0] for y in ys]) # <---- This line
>> print(ndx)
[1 4]

Есть ли быстрый способ?

Спасибо

Answer 1

Для больших массивов xs и ys вам нужно изменить базовый подход, чтобы это стало быстрым. Если у вас все в порядке с сортировкой xs, то простой вариант - использовать numpy.searchsorted():

xs.sort()
ndx = numpy.searchsorted(xs, ys)

Если важно сохранить исходный порядок xs, вы также можете использовать этот подход, но вам нужно запомнить исходные индексы:

orig_indices = xs.argsort()
ndx = orig_indices[numpy.searchsorted(xs[orig_indices], ys)]