将多个数组与 numpy 数组相交的最佳方法是什么？

Question

假设我有一个numpy数组的例子：

import numpy as np
X = np.array([2,5,0,4,3,1])

我还有一个数组列表，例如：

A = [np.array([-2,0,2]), np.array([0,1,2,3,4,5]), np.array([2,5,4,6])]

我只想保留每个列表中也在 X 中的这些项目。我也希望以最有效/最常见的方式做到这一点。

到目前为止我尝试过的解决方案：

1. X使用排序X.sort()。

2. 使用以下方法查找每个数组的项目位置X：

   locations = [np.searchsorted(X, n) for n in A]
   

3. 只留下合适的：

   masks = [X[locations[i]] == A[i] for i in range(len(A))]
   result = [A[i][masks[i]] for i in range(len(A))]
   

但它不起作用，因为第三个数组的位置超出范围：

locations = [array([0, 0, 2], dtype=int64), array([0, 1, 2, 3, 4, 5], dtype=int64), array([2, 5, 4, 6], dtype=int64)]

如何解决这个问题？

更新

我最终得到了idx[idx==len(Xs)] = 0解决方案。我还注意到答案之间发布了两种不同的方法：转换X为setvs np.sort。它们都有优点和缺点：set操作使用迭代，与方法相比相当慢；numpy然而，与即时np.searchsorted访问项目不同，速度呈对数增长。set这就是为什么我决定使用大尺寸数据（尤其是 100 万个项目）来比较性能X, A[0], A[1], A[2]。

Answer 1

一种想法是循环时减少计算和最少的工作。所以，这是一个考虑到这些的——

a = np.concatenate(A)
m = np.isin(a,X)
l = np.array(list(map(len,A)))
a_m = a[m]
cut_idx = np.r_[0,l.cumsum()]
l_m = np.add.reduceat(m,cut_idx[:-1])
cl_m = np.r_[0,l_m.cumsum()]
out = [a_m[i:j] for (i,j) in zip(cl_m[:-1],cl_m[1:])]

替代方案#1：

我们还可以用来np.searchsorted获取isin面具，就像这样 -

Xs = np.sort(X)
idx = np.searchsorted(Xs,a)
idx[idx==len(Xs)] = 0
m = Xs[idx]==a

另一种方式是np.intersect1d

如果您正在寻找最常见/优雅的一种，请认为它是np.intersect1d-

In [43]: [np.intersect1d(X,A_i) for A_i in A]
Out[43]: [array([0, 2]), array([0, 1, 2, 3, 4, 5]), array([2, 4, 5])]

解决您的问题

您还可以通过简单的修复来解决越界问题 -

for l in locations:
    l[l==len(X)]=0