Dave Dave - 8 months ago 34
Python Question

Return common element indices between two numpy arrays

I have two arrays, a1 and a2. Assume

len(a2) >> len(a1)
, and that a1 is a subset of a2.

I would like a quick way to return the a2 indices of all elements in a1. The time-intensive way to do this is obviously:

from operator import indexOf
indices = []
for i in a1:

This of course takes a long time where a2 is large. I could also use numpy.where() instead (although each entry in a1 will appear just once in a2), but I'm not convinced it will be quicker. I could also traverse the large array just once:

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

But I'm sure there is a faster, more 'numpy' way - I've looked through the numpy method list, but cannot find anything appropriate.

Many thanks in advance,



How about

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

This should be fast. From my basic testing, it's about 7 times faster than your second code snippet for len(a2) == 100, len(a1) == 10000, and only one common element at index 45. This assumes that both a1 and a2 have no repeating elements.