Manera eficiente de calcular valores de intersección entre dos matrices numpy

Tengo un cuello de botella en mi progtwig que es causado por lo siguiente:

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

El resultado esperado para C es el siguiente:

 C = [4 6 7 1 5 4 1 1 9] 

¿Hay una forma más eficiente de hacer esta operación?

Tenga en cuenta que la matriz A contiene valores que se repiten y deben tenerse en cuenta. No pude usar la intersección de conjuntos ya que al tomar la intersección se omitirán los valores de repetición, devolviendo solo [1,4,5,6,7,9] .

También tenga en cuenta que esto es sólo una simple demostración. Los tamaños de matriz reales pueden ser del orden de miles, hasta más de millones.

Puedes usar np.in1d :

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

np.in1d devuelve una matriz booleana que indica si cada valor de A también aparece en B Esta matriz se puede usar para indexar A y devolver los valores comunes.

No es relevante para su ejemplo, pero también vale la pena mencionar que si A y B contienen valores únicos, entonces se puede acelerar assume_unique=True configurando assume_unique=True :

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

También puede interesarle np.intersect1d que devuelve una matriz de los valores únicos comunes a ambas matrices (ordenados por valor):

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

Utilice numpy.in1d :

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

Si solo verificas la existencia en B ( if i in B ) entonces obviamente puedes usar un set para esto. No importa cuántos cuatros haya en B , siempre que haya al menos uno. Por supuesto que tienes razón, que no puedes usar dos conjuntos y una intersección. Pero incluso un set debería mejorar el rendimiento, ya que la complejidad de búsqueda es menor que 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]) 

Podemos usar np.searchsorted para boost el rendimiento, más aún en el caso de que la matriz de búsqueda haya ordenado valores únicos:

 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] 

Esa bandera assume_unique hace que funcione tanto para el caso genérico como para el caso especial de que B sea ​​único y esté ordenado.

Ejecución de la muestra

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

Tiempos para comparar con otra solución vectorizada basada en np.in1d (enumerada en otras dos respuestas) en arreglos grandes para ambos casos:

 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