Каков наилучший способ пересечения нескольких массивов с массивом 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 в set против np.sort. У обоих из них есть свои плюсы и минусы: set операции используют итерации, которые довольно медленные по сравнению с numpy методами; однако np.searchsorted скорость увеличивается логарифмически в отличие от принятия set предметов, что мгновенно. Поэтому я решил сравнить производительность, используя данные огромных размеров, особенно 1 миллион элементов для 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

Answer 2

Как насчет этого, очень просто и эффективно:

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)