计算两个 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].

另请注意，这只是一个简单的演示。实际的数组大小可以是数千个，甚至超过数百万个。

Answer 1

使用numpy.in1d:

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

Answer 2

您可以使用 np.in1d:

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

np.in1d returns 一个布尔数组，指示 A 的每个值是否也出现在 B 中。然后可以使用此数组索引 A 和 return 常用值。

这与您的示例无关，但也值得一提的是，如果 A 和 B 每个都包含唯一值，则 np.in1d 可以通过设置 assume_unique=True:

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

您可能还对 np.intersect1d 感兴趣，其中 return 是两个数组共有的唯一值数组（按值排序）：

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

Answer 3

如果您只检查 B (if i in B) 中是否存在，那么显然您可以为此使用 set。只要至少有一个，B 中有多少个四边形并不重要。当然你是对的，你不能使用两个集合和一个交集。但即使是一个 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 4

我们可以使用 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

Answer 5

1-使用设置的交集，因为在这种情况下它非常快

c = set(a).intersection(b)

2-使用numpy intersect1d方法，比循环快但比第一种方法慢

c = numpy.intersect1d(a,b)

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

Efficient way to compute intersecting values between two numpy arrays

python

arrays

numpy