python - 将 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”有问题。我尝试直接使用 bool 值 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