计算两个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
并返回公共值。
这与您的示例无关,但值得一提的是,如果A
和B
每个都包含唯一值,则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])