element wise “contains” in Python

Say I have an array:

import numpy as np
arr = np.random.randint(0, 5, 20)

then arr>3 results in an array of type bool with shape (20,). How can I most efficiently do the same thing with the “contains” operator? The simple

arr in [2, 4] 

will result in “The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()”. Is there another way than

np.array([ x in [2, 4] for x in arr])

?

Answer

You can use np.isin:

import numpy as np
arr = np.random.randint(0, 5, 20)
np.isin(arr, [2, 4])

Output:

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

The function returns a boolean array of the same shape as your input array named arr that is True where an element of arr is in your second list argument [2, 4] and False otherwise