我的程序中存在瓶颈,原因如下：

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])