查找另一个数组中一个数组的匹配索引

Question

我有两个numpy数组,A和B.一个包含唯一值,B是A的子数组.现在我正在寻找一种方法来获得B中B值的索引.

例如:

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

Answer 1

您可以使用np.in1d带有np.nonzero-

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

np.searchsorted如果您关心维护订单,您也可以使用-

np.searchsorted(A,B)

对于一般情况,当A&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])

Answer 2

你试过了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])