numpy conditional list membership element wise

Question

I have a 2D numpy array:

a = np.array([[0,1],
              [2,3]])

I have a list of values to keep:

vals_keep = [1,2]

I want to test for list membership for each element in the array. Something like:

mask = a in vals_keep

The result I want:

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

Answer 1

You can use isin

isin is an element-wise function version of the python keyword in

np.isin(a, vals_keep)

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

An added benefit of isin is that it's flexible with arrays of different dimensions:

a = np.arange(4).reshape(1,2,2,1)
np.isin(a, vals_keep)

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

        [[ True],
         [False]]]])

Answer 2

Here is one way using broadcasting:

In [35]: (a[:, :, None] == vals_keep).any(2)
Out[35]: 
array([[False,  True],
       [ True, False]])

Which is faster than isin for small arrays (less than 100 rows):

In [37]: %timeit np.isin(a, vals_keep)
22 µs ± 728 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [38]: %timeit (a[:, :, None] == vals_keep).any(2)
12.6 µs ± 95.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

For large arrays it's better to use isin because broadcasting in 3D is not very efficient for large arrays/matrices.