Search code examples
pythonarraysnumpymembership

Mask out specific values from an array

Example:

I have an array:

array([[1, 2, 0, 3, 4],
       [0, 4, 2, 1, 3],
       [4, 3, 2, 0, 1],
       [4, 2, 3, 0, 1],
       [1, 0, 2, 3, 4],
       [4, 3, 2, 0, 1]], dtype=int64)

I have a set (variable length, order doesn't matter) of "bad" values:

{2, 3}

I want to return the mask that hides these values:

array([[False,  True, False,  True, False],
       [False, False,  True, False,  True],
       [False,  True,  True, False, False],
       [False,  True,  True, False, False],
       [False, False,  True,  True, False],
       [False,  True,  True, False, False]], dtype=bool)

What's the simplest way to do this in NumPy?

Solution

  • Use np.in1d that gives us a flattened mask of such matching occurrences and then reshape back to input array shape for the desired output, like so -

    np.in1d(a,[2,3]).reshape(a.shape)

    Note that we need to feed in the numbers to be searched as a list or an array.

    Sample run -

    In [5]: a
Out[5]: 
array([[1, 2, 0, 3, 4],
       [0, 4, 2, 1, 3],
       [4, 3, 2, 0, 1],
       [4, 2, 3, 0, 1],
       [1, 0, 2, 3, 4],
       [4, 3, 2, 0, 1]])

In [6]: np.in1d(a,[2,3]).reshape(a.shape)
Out[6]: 
array([[False,  True, False,  True, False],
       [False, False,  True, False,  True],
       [False,  True,  True, False, False],
       [False,  True,  True, False, False],
       [False, False,  True,  True, False],
       [False,  True,  True, False, False]], dtype=bool)

    2018 Edition : numpy.isin

    Use NumPy built-in np.isin (introduced in 1.13.0) that keeps the shape and hence doesn't require us to reshape afterwards -

    In [153]: np.isin(a,[2,3])
Out[153]: 
array([[False,  True, False,  True, False],
       [False, False,  True, False,  True],
       [False,  True,  True, False, False],
       [False,  True,  True, False, False],
       [False, False,  True,  True, False],
       [False,  True,  True, False, False]])