Masking a 2D numpy array by comparing elements to a 1D list

September 11, 2022

Say I’m given a 2D numpy array and a list of numbers that look like this:

[1 2 3 1]
[2 3 1 2]
[3 1 2 3]
[1 2 3 1]

[1, 3]

I need to find some way to generate a boolean array where every element that’s in the list results in a True in a 2D array

[True False True True]
[False True True False]
[True True False True]
[True False True True]

I’ve tried looping through the list and using logical_or to combine the boolean masks that it creates, but the 2D array is a large image and the list can be upwards of 300 elements long, which becomes very slow very quickly.

If there’s some operation in numpy that looks like array in list similar to array == list[0], I can’t find it but it would be very helpful.

My code currently looks like

boolean_array = array == mask_list[0]
for list_element in mask_list[0:]:
    boolean_array = np.logical_or(blank_boolean, array == list_element)

return boolean_array

(np.where, np.any, np.all have also been tried but i cant wrap my head around how to do it)

(also both sets are always integers)

>Solution :

Use np.isin

import numpy as np

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

valid = np.array([1, 3])

print(np.isin(arr, valid))

[[ True False  True  True]
 [False  True  True False]
 [ True  True False  True]
 [ True False  True  True]]