两个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 包含重复的值，并且它们需要被考虑在内。我无法使用集合交集，因为取交集会省略重复的值，只返回 [1,4,5,6,7,9]。

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

- user32147

Answer 1

您可以使用np.in1d：

>>> 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感兴趣，它返回一个数组，其中包含两个数组共有的唯一值（按值排序）：

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

- Alex Riley

---

假设两个数组是唯一的，我们可以使用np.in1d和np.intersect1d之一。您能否评论一下它们之间的性能差异？ - user32147

我还没有对这两个方法的性能进行详尽的测试，但是如果在两种方法中都将assume_unique设置为True，则np.intersect1d似乎稍微更快一些。我不确定确切的原因是什么，但可能是因为它需要进行较少的比较。 - Alex Riley

Answer 2

使用 numpy.in1d 方法：

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

- Ashwini Chaudhary

Answer 3

我们可以使用 np.searchsorted 来提高性能，特别是在查找数组具有排序唯一值的情况下更为明显。

def intersect1d_searchsorted(A,B,assume_unique=False):
    if assume_unique==0:
        B_ar = np.unique(B)
    else:
        B_ar = B
    idx = np.searchsorted(B_ar,A)
    idx[idx==len(B_ar)] = 0
    return A[B_ar[idx] == A]

那个assume_unique标志使其适用于通用情况和B是唯一且已排序的特殊情况。

示例运行 -

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

In [90]: intersect1d_searchsorted(A,B,assume_unique=True)
Out[90]: array([4, 6, 7, 1, 5, 4, 1, 1, 9])

在两种情况下，与另一个基于向量化的np.in1d解决方案比较大的数组的时间 -

In [103]: A = np.random.randint(0,10000,(1000000))

In [104]: B = np.random.randint(0,10000,(1000000))

In [105]: %timeit A[np.in1d(A, B)]
     ...: %timeit A[np.in1d(A, B, assume_unique=False)]
     ...: %timeit intersect1d_searchsorted(A,B,assume_unique=False)
1 loop, best of 3: 197 ms per loop
10 loops, best of 3: 190 ms per loop
10 loops, best of 3: 151 ms per loop

In [106]: B = np.unique(np.random.randint(0,10000,(5000)))

In [107]: %timeit A[np.in1d(A, B)]
     ...: %timeit A[np.in1d(A, B, assume_unique=True)]
     ...: %timeit intersect1d_searchsorted(A,B,assume_unique=True)
10 loops, best of 3: 130 ms per loop
1 loop, best of 3: 218 ms per loop
10 loops, best of 3: 80.2 ms per loop

- Divakar

Answer 4

1-在这种情况下，请使用集合交集，因为它非常快

c = set(a).intersection(b)

2-使用NumPy的intersect1d方法，它比循环更快但比第一个方法慢。

c = numpy.intersect1d(a,b)

- khaled khaled

Answer 5

如果您只是检查B中是否存在某个元素（if i in B），那么您显然可以使用set来实现。无论B中有多少个“four”，只要至少有一个即可。当然，您是正确的，不能使用两个集合和一个交集。但是，即使只使用一个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])

- BartoszKP