计算两个numpy数组之间的相交值的有效方法


问题内容

我的程序中存在瓶颈,这是由以下原因引起的:

A = numpy.array([10,4,6,7,1,5,3,4,24,1,1,9,10,10,18])
B = numpy.array([1,4,5,6,7,8,9])

C = numpy.array([i for i in A if i in B])

的预期结果C如下:

C = [4 6 7 1 5 4 1 1 9]

有没有更有效的方法来执行此操作?

请注意,数组A包含重复值,需要将它们考虑在内。我无法使用设置交集,因为采用交集会忽略重复值,而返回just [1,4,5,6,7,9]

还要注意,这只是一个简单的演示。实际的数组大小可能在数千个级别,甚至超过数百万个。


问题答案:

您可以使用np.in1d

>>> A[np.in1d(A, B)]
array([4, 6, 7, 1, 5, 4, 1, 1, 9])

np.in1d返回一个布尔数组,指示的每个值是否A也出现在中B。然后,可以使用该数组索引A并返回公共值。

这与您的示例无关,但值得一提的是,如果AB每个都包含唯一值,则np.in1d可以通过设置来加快速度assume_unique=True

np.in1d(A, B, assume_unique=True)

您可能还对以下np.intersect1d哪项感兴趣:返回两个数组共有的唯一值的数组(按值排序):

>>> np.intersect1d(A, B)
array([1, 4, 5, 6, 7, 9])