searchsorted
如果不是 B 的每个元素都在 A 中,可能会给您错误的答案。您可以使用numpy.in1d
:
A = np.array([1,2,3,4,5,6,7])
B = np.array([2,4,6,8])
mask = np.in1d(A, B)
print np.where(mask)[0]
print np.where(~mask)[0]
输出是:
[1 3 5]
[0 2 4 6]
但是in1d()
使用排序,这对于大型数据集来说很慢。如果您的数据集很大,您可以使用 pandas:
import pandas as pd
np.where(pd.Index(pd.unique(B)).get_indexer(A) >= 0)[0]
下面是时间对比:
A = np.random.randint(0, 1000, 10000)
B = np.random.randint(0, 1000, 10000)
%timeit np.where(np.in1d(A, B))[0]
%timeit np.where(pd.Index(pd.unique(B)).get_indexer(A) >= 0)[0]
输出:
100 loops, best of 3: 2.09 ms per loop
1000 loops, best of 3: 594 µs per loop