python - 将多个数组与 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 中的这些项目。我也希望以最有效/最常见的方式来做。

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

排序 X使用 X.sort() .

在 X 中查找每个数组的项的位置使用:

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

只留下适当的:

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对比 np.sort .它们都有优点和缺点: set操作使用迭代，即 quite slow与 numpy 相比方法;然而 np.searchsorted与访问 set 不同，速度以对数方式增加即时的项目。这就是为什么我决定使用大尺寸数据来比较性能，尤其是 X, A[0], A[1], A[2] 的 100 万个项目。 .

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])]

根据 docs ，集合操作都有 O(1)复杂性，因此总体为 O(N)