使用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. Sort X using X.sort().

2. Find locations of items of each array in X using:

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

3. Leave only proper ones:

   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转换为set vs np.sort。 它们都有优缺点：set操作使用迭代，与numpy方法相比非常慢; 但是np.searchsorted的速度呈对数增长，不像访问set项一样快速。 这就是为什么我决定使用具有巨大大小的数据进行性能比较，特别是对于X，A[0]，A[1]，A[2]拥有100万个项目的数据。

- mathfux

Answer 1

这个怎么样，非常简单高效：

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_set = set(X)
A = [np.array([a for a in arr if a in X_set]) for arr in A]
#[array([0, 2]), array([0, 1, 2, 3, 4, 5]), array([2, 5, 4])]

根据文档，集合操作的复杂度均为O(1)，因此总复杂度为O(N)。

- ExplodingGayFish

Answer 2

一个思路是在循环时减少计算量和最小化工作量。因此，以下代码考虑了这些要点 -

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

- Divakar

---

我相信使用intersect1d是我的想法的一个很好的表达，但我处理的X和A非常巨大，所以我决定先对X进行排序。 - mathfux

@mathfux 我认为你会发现，对于非常大的数组来说，排序比使用set函数要慢得多。 - Daniel F