Filter A N-d Numpy Array And Keep Only Specific Elements

I'm dealing with a large N-D numpy array. I would like to keep only those elements present in a different numpy array, and set the remaining values to 0. for example, if we conside

Solution 1:

If you are onto using only numpy, this can also be done using simple use of broadcasting by casting the vals array to just one rank higher than a. This is accomplished without using iterations or other functionalities.

import numpy as np

a = np.array([[[36,  1, 72],
         [76, 50, 23],
         [28, 68, 17],
         [84, 75, 69]],
 
        [[ 5, 15, 93],
         [92, 92, 88],
         [11, 54, 21],
         [87, 76, 81]]])

vals = np.array([50, 11, 72])
inds = a == vals[:, None, None, None]
a[~np.any(inds, axis = 0)] = 0
a

Output:

array([[[ 0,  0, 72],
        [ 0, 50,  0],
        [ 0,  0,  0],
        [ 0,  0,  0]],

       [[ 0,  0,  0],
        [ 0,  0,  0],
        [11,  0,  0],
        [ 0,  0,  0]]])

Solution 2:

I set up a mask by combining reduce with np.logical_or and iterated over the values that should remain:

import functools
import numpy as np

arr = np.array([[[36,  1, 72],
        [76, 50, 23],
        [28, 68, 17],
        [84, 75, 69]],
       [[ 5, 15, 93],
        [92, 92, 88],
        [11, 54, 21],
        [87, 76, 81]]])

# Set the values that should not
# be set to zero
vals = [11, 50, 72]

# Create a mask by looping over the above values
mask = functools.reduce(np.logical_or, (arr==val for val in vals))

masked = np.where(mask, arr, 0.)

print(masked)
> array([[[ 0.,  0., 72.],
        [ 0., 50.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [11.,  0.,  0.],
        [ 0.,  0.,  0.]]])