我的程序中存在瓶颈,原因如下:
A = numpy.array([10,4,6,7,1,5,3,24,9,10,18]) B = numpy.array([1,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包含重复值,需要将它们考虑在内.我无法使用集合交集,因为取交点将省略重复值,仅返回[1,9].
另请注意,这只是一个简单的演示.实际的阵列大小可以是数千,而不是数百万.
解决方法
你可以使用
np.in1d
:
>>> A[np.in1d(A,B)] array([4,9])
np.in1d返回一个布尔数组,指示A的每个值是否也出现在B.这个数组然后可用于索引A并返回公共值.
它与你的例子无关,但是值得一提的是,如果A和B都包含唯一值,那么可以通过设置assume_unique = True来加速np.in1d:
np.in1d(A,B,assume_unique=True)
您可能还对np.intersect1d
感兴趣,它返回两个数组共有的唯一值数组(按值排序):
>>> np.intersect1d(A,B) array([1,9])