Find indices of large array if it contains values in smaller array

Question

Is there a fast numpy function for returning a list of indices in a larger array where it matches values from a smaller array? Smaller array is ~ 30M values and larger is 800M so I want to avoid a for-loop of numpy.where calls.

The problem with searchsorted is that it will return results even when their is not an exact match, it just gives the closest index, but I only want indices where there are exact matches

instead of this:

>>> a = array([1,2,3,4,5])
>>> b = array([2,4,7])
>>> searchsorted(a,b)
array([1, 3, 5])

I would want this:

>>> a = array([1,2,3,4,5])
>>> b = array([2,4,7])
>>> SOMEFUNCTION(a,b)
array([1, 3])

EDIT: the set of values in both the smaller and larger arrays are always unique and sorted.

Answer 1

You could use np.in1d to find those elements of a which are in b. To find the index, use one call to np.where:

In [34]: a = array([1,2,3,4,5])

In [35]: b = array([2,4,7])

In [36]: np.in1d(a, b)
Out[38]: array([False,  True, False,  True, False], dtype=bool)

In [39]: np.where(np.in1d(a, b))
Out[39]: (array([1, 3]),)

Because a and b are already sorted, you could use

In [57]: np.searchsorted(b, a, side='right') != np.searchsorted(b, a, side='left')
Out[57]: array([False,  True, False,  True, False], dtype=bool)

instead of np.in1d(a, b). For large a and b, using searchsorted may be faster:

import numpy as np
a = np.random.choice(10**7, size=10**6, replace=False)
a.sort()
b = np.random.choice(10**7, size=10**5, replace=False)
b.sort()

In [53]: %timeit np.in1d(a, b)
10 loops, best of 3: 176 ms per loop

In [54]: %timeit np.searchsorted(b, a, side='right') != np.searchsorted(b, a, side='left')
10 loops, best of 3: 106 ms per loop

---

Jaime and Divakar have suggested some significant improvements on the method shown above. Here is some code which tests that the methods all return the same result, followed by some benchmarks:

import numpy as np

a = np.random.choice(10**7, size=10**6, replace=False)
a.sort()
b = np.random.choice(10**7, size=10**5, replace=False)
b.sort()

def using_searchsorted(a, b):
    return (np.where(np.searchsorted(b, a, side='right') 
                     != np.searchsorted(b, a, side='left')))[0]

def using_in1d(a, b):
    return np.where(np.in1d(a, b))[0]

def using_searchsorted_divakar(a, b):
    idx1 = np.searchsorted(a,b,'left')
    idx2 = np.searchsorted(a,b,'right')
    out = idx1[idx1 != idx2]
    return out

def using_jaime_mask(haystack, needle):
    idx = np.searchsorted(haystack, needle)
    mask = idx < haystack.size
    mask[mask] = haystack[idx[mask]] == needle[mask]
    idx = idx[mask]
    return idx

expected = using_searchsorted(a, b)
for func in (using_in1d, using_searchsorted_divakar, using_jaime_mask):
    result = func(a, b)
    assert np.allclose(expected, result)

---

In [29]: %timeit using_jaime_mask(a, b)
100 loops, best of 3: 13 ms per loop

In [28]: %timeit using_searchsorted_divakar(a, b)
10 loops, best of 3: 21.7 ms per loop

In [26]: %timeit using_searchsorted(a, b)
10 loops, best of 3: 109 ms per loop

In [27]: %timeit using_in1d(a, b)
10 loops, best of 3: 173 ms per loop

Answer 2

The default seacrh direction with np.searchsorted is left. We can also search for it from the right direction and the ones that are same in both these set of indices would be the ones to be avoided from the indices outputted from the left option to get the desired output. The motivation here is same as discussed in @unutbu's solution.

Thus, the implementation would look like this -

idx1 = np.searchsorted(a,b,'left')
idx2 = np.searchsorted(a,b,'right')
out = idx1[idx1 != idx2]