使用numpy.in1d()时，数组与单个元素的结果不同

Question

我正在用Python写代码，但遇到了一些问题。我有两个数组，假设叫做A和B，里面都是一些ID。A包含了所有的ID，而B则是属于某个组的ID。我想做的是找出B中的元素在A中的位置，使用的代码是：

>>> print B
[11600813 11600877 11600941 ..., 13432165 13432229 13434277]
>>> mask=np.nonzero(np.in1d(A, B))
>>> print A[mask]
[12966245 12993389 12665837 ..., 13091877 12965029 13091813]

但是这显然不对，因为我没有得到B的值。我检查了一下自己是否正确使用了numpy.in1d()，于是我尝试了：

>>> mask=np.nonzero(np.in1d(A, B[0]))
>>> print A[mask]
[11600813]

这个是对的，所以我猜测在numpy.in1d()中‘B’可能有问题。我还尝试直接使用布尔值np.in1d(A, B)，而不是把它转换成索引，但也没有成功。我还试过用B = numpy.array(B)和B = list(B)，但都不行。

不过如果我用B = numpy.array(B)[0]或者B = list(B)[0]，对于那个元素来说是可以的。可惜的是，我不能对每个元素都用一个'for'循环，因为len(A)是16777216，而len(B)是9166，这样会花费很多时间。

我还确保了B中的所有元素都在A中：

>>> np.intersect1d(A, B)
[11600813 11600877 11600941 ..., 13432165 13432229 13434277]

Answer 1

你可以使用 numpy.argsort 和 numpy.searchsorted 来获取位置：

import numpy as np
A = np.unique(np.random.randint(0, 100, 100))
B = np.random.choice(A, 10)

idxA = np.argsort(A)
sortedA = A[idxA]
idxB = np.searchsorted(sortedA, B)
pos = idxA[idxB]
print A[pos]
print B

如果你想要更快的方法，可以考虑使用 pandas。

import pandas as pd
s = pd.Index(A)
pos = s.get_indexer(B)
print A[pos]
print B