Mask of boolean 2D numpy array with True values for elements contained in another 1D numpy array-CodePudding

Take the following example. I have an array test and want to get a boolean mask with True's for all elements that are equal to elements of ref.

import numpy as np
test = np.array([[2, 3, 1, 0], [5, 4, 2, 3], [6, 7, 5 ,4]])
ref = np.array([3, 4, 5])

I am looking for something equivalent to

mask = (test == ref[0]) | (test == ref[1]) | (test == ref[2])

which in this case should yield

>>> print(mask)
[[False,  True, False, False],
 [ True,  True, False,  True],
 [False, False,  True,  True]]

but without having to resort to any loops.

CodePudding user response：

Numpy comes with a function isin that does exactly this

np.isin(test, ref)

which return

array([[False,  True, False, False],
       [ True,  True, False,  True],
       [False, False,  True,  True]])

CodePudding user response：

You can use numpy broadcasting:

mask = (test[:,None] == ref[:,None]).any(1)

output:

array([[False,  True, False, False],
       [ True,  True, False,  True],
       [False, False,  True,  True]])

NB. this is faster that numpy.isin, but creates a (X, X, Y) sized intermediate array where X, Y is the shape of test, so this will consume some memory on very large arrays