Ottenere gli indici di più elementi in una matrice NumPy allo stesso tempo

Question

C'è un modo per ottenere gli indici di più elementi in una matrice NumPy in una volta?

E.g.

import numpy as np 
a = np.array([1, 2, 4]) 
b = np.array([1, 2, 3, 10, 4]) 

vorrei trovare l'indice di ogni elemento di a in b, vale a dire: [0,1,4].

trovo la soluzione che sto usando un po 'prolisso:

import numpy as np 

a = np.array([1, 2, 4]) 
b = np.array([1, 2, 3, 10, 4]) 

c = np.zeros_like(a) 
for i, aa in np.ndenumerate(a): 
    c[i] = np.where(b==aa)[0] 

print('c: {0}'.format(c)) 

uscita:

c: [0 1 4] 

Answer 1

Si potrebbe utilizzare in1d e nonzero (o where per questo):

>>> np.in1d(b, a).nonzero()[0] 
array([0, 1, 4]) 

Questo funziona bene per gli array di esempio, ma in generale l'ar il raggio degli indici restituiti non rispetta l'ordine dei valori in a. Questo potrebbe essere un problema a seconda di cosa vuoi fare dopo.

In questo caso, una risposta molto migliore è quella @Jaime dà here, utilizzando searchsorted:

>>> sorter = np.argsort(b) 
>>> sorter[np.searchsorted(b, a, sorter=sorter)] 
array([0, 1, 4]) 

Ciò restituisce gli indici per i valori come appaiono in a. Per esempio:

a = np.array([1, 2, 4]) 
b = np.array([4, 2, 3, 1]) 

>>> sorter = np.argsort(b) 
>>> sorter[np.searchsorted(b, a, sorter=sorter)] 
array([3, 1, 0]) # the other method would return [0, 1, 3] 

Answer 2

Questo è un semplice one-liner utilizzando il pacchetto numpy-indexed (disclaimer: io sono il suo autore):

import numpy_indexed as npi 
idx = npi.indices(b, a) 

L'implementazione è pienamente vettorializzare, e consente di controllare la gestione dei valori mancanti. Inoltre, funziona anche per gli array nd (ad esempio, trovando gli indici delle righe di a in b).