Search code examples
numpysearchfilternumba

Search multiple elements in 2D array?

I have a 1d array match and 2d array data . I can search elements one or several like this :

 data[data[:,0] == match[0] ]
 data[data[:,0] == match[0] & data[:,0] == match[1] ]

but how do you search for all of them :

 data[ data[:,0] == match ]

I suppose inside numba function i can use from in1d() !! :

    mask = np.zeros(len(ar1), dtype=bool)
    for a in ar2:
            mask |= (ar1 == a)
    return mask
Solution

  • You can use numpy.isin:

    np.isin(data[:, 0], match)

    Example:

    data = np.array([[1, 4],[5, 2],[2, 4]])
match = np.array([2, 4])

np.isin(data[:,0], match)
# array([False, False,  True])

data[np.isin(data[:, 0], match)]
# array([[2, 4]])