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

Question

我的程序遇到瓶颈，原因如下：

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包含重复值，需要考虑这些值。我无法使用set intersection，因为获取intersection将忽略重复值，只返回[1,4,5,6,7,9]。

还要注意，这只是一个简单的演示。实际的数组大小可以是数千到数百万的顺序。

Answer 1

如果您只检查B（if i in B）中是否存在，那么显然您可以使用set来完成此操作。只要至少有一个，在B中有多少个fours并不重要。当然你是对的，你不能用两个集合和一个交集。但即使是一个set也应该提高性能，因为搜索复杂性小于O（n）：

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

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

Answer 2

使用^{}：

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

Answer 3

您可以使用^{}：

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

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

这与您的示例无关，但也值得一提的是，如果A和B各自包含唯一值，则可以通过设置assume_unique=True来加快np.in1d：

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

您可能还对^{}感兴趣，它返回两个数组通用的唯一值数组（按值排序）：

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