Вы можете использовать np.in1d
с np.nonzero
-
np.nonzero(np.in1d(A,B))[0]
Вы также можете использовать np.searchsorted
, если вы хотите сохранить порядок -
np.searchsorted(A,B)
Для общего случая, когда A
& amp; B
- несортированные массивы, вы можете ввести параметр sorter
в 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])