2010-02-25 30 views
9

Ho due matrici, a1 e a2. Si assuma len(a2) >> len(a1) e che a1 sia un sottoinsieme di a2.Restituisce indici di elementi comuni tra due array numpy

Vorrei un modo rapido per restituire gli indici a2 di tutti gli elementi in a1. Il modo più intenso per farlo è ovviamente:

from operator import indexOf 
indices = [] 
for i in a1: 
    indices.append(indexOf(a2,i)) 

Questo ovviamente richiede molto tempo quando a2 è grande. Potrei anche usare numpy.where() (anche se ogni voce in a1 apparirà solo una volta in a2), ma non sono convinto che sarà più veloce. Potrei anche attraversare la vasta gamma solo una volta:

for i in xrange(len(a2)): 
    if a2[i] in a1: 
     indices.append(i) 

Ma sono sicuro che ci sia un modo più veloce e 'NumPy' - Ho guardato l'elenco metodo NumPy, ma non riesco a trovare nulla appropriata.

Molte grazie in anticipo,

D

risposta

8

Come su

numpy.nonzero(numpy.in1d(a2, a1))[0] 

Questo dovrebbe essere veloce. Dal mio test di base, è circa 7 volte più veloce del tuo secondo snippet di codice per len(a2) == 100, len(a1) == 10000 e solo un elemento comune all'indice 45. Ciò presuppone che sia a1 e a2 non abbiano elementi ripetuti.

+0

Ho confrontato la tua soluzione con quella di Dave Kirby di cui sopra, con questo a circa 1.35X più veloce per len (a2) == 12347424, len (a1) == 1338, quindi questa soluzione è il mio voto - grazie! – Dave

+3

Per chi legge questo: sembra che 'setmember1d' sia stato rinominato in' in1d' da numpy 1.4. –

+0

@AlokSinghal Grazie per averlo preso! –

2

come circa:

wanted = set(a1) 
indices =[idx for (idx, value) in enumerate(a2) if value in wanted] 

Questo dovrebbe essere O (len (A1) + len (a2)) invece di O (len (a1) * len (a2))

NB Non so numpy quindi potrebbe esserci un modo più 'numpythonic' per farlo, ma questo è come lo farei in puro python.

+0

dovrebbe essere elencato (a2)? – Dave

+0

Oops, mio ​​male. Risolto adesso –

1
index = in1d(a2,a1) 
result = a2[index] 
0

Il pacchetto numpy_indexed (disclaimer: io sono il suo autore) contiene un equivalente di vettorializzare indice di una lista; la performance dovrebbe essere simile alla risposta attualmente accettata, ma come bonus, ti dà anche il controllo esplicito sui valori mancanti, usando il kwarg 'mancante'.

import numpy_indexed as npi 
indices = npi.indices(a2, a1, missing='raise') 

Inoltre, funzionerà anche sugli array multidimensionali, cioè, trovando gli indici di una serie di righe di un'altra.

1

Molto simile a @AlokSinghal, ma si ottiene una versione già appiattita.

numpy.flatnonzero(numpy.in1d(a2, a1)) 
Problemi correlati