Поиск индексов совпадений одного массива в другом массиве

У меня есть два массива numpy, A и B. A conatains уникальные значения, А B-поддерево A. Теперь я ищу способ получить индекс значений B в пределах A.

например:

A = np.array([1,2,3,4,5,6,7,8,9,10])
B = np.array([1,7,10])
# I need a function fun() that:
fun(A,B)
>> 0,6,9

3 ответов


можно использовать np.in1d С np.nonzero -

np.nonzero(np.in1d(A,B))[0]

вы также можете использовать np.searchsorted, если вы заботитесь о поддержании порядка -

np.searchsorted(A,B)

для общего случая, когда A & B являются несортированными массивами, вы можете ввести на np.searchsorted, вот так -

sort_idx = A.argsort()
out = sort_idx[np.searchsorted(A,B,sorter = sort_idx)]

Я бы добавил в моем любимом broadcasting тоже в смеси, чтобы решить общий случай -

np.nonzero(B[:,None] == A)[1]

образец выполнения -

In [125]: A
Out[125]: array([ 7,  5,  1,  6, 10,  9,  8])

In [126]: B
Out[126]: array([ 1, 10,  7])

In [127]: sort_idx = A.argsort()

In [128]: sort_idx[np.searchsorted(A,B,sorter = sort_idx)]
Out[128]: array([2, 4, 0])

In [129]: np.nonzero(B[:,None] == A)[1]
Out[129]: array([2, 4, 0])

ты пробовал searchsorted?

A = np.array([1,2,3,4,5,6,7,8,9,10])
B = np.array([1,7,10])

A.searchsorted(B)
# array([0, 6, 9])

просто для полноты: если значения в A не отрицательны и достаточно малы:

lookup = np.empty((np.max(A) + 1), dtype=int)
lookup[A] = np.arange(len(A))
indices  = lookup[B]