Index numpy arrays with a list of numbers

Question

I want to index a numpy array based on specific numbers, eg, I have the following np array,

b = np.array([[1, 2, 3, 2, 1, 3],[2, 1, 3, 5, 4, 1]])

And I want to change to zero the values different from 1 and 2. I tried,

b[b not in [1, 2]] = 0

but did work. Any ideas on how to do it?

Answer 1

You can use numpy.isin() :

import numpy as np
b = np.array([[1, 2, 3, 2, 1, 3],[2, 1, 3, 5, 4, 1]])
v = np.array([1,2])
m = np.isin(b, v)
b[~m] = 0
print(b)

Output:

[[1 2 0 2 1 0]
 [2 1 0 0 0 1]]

Answer 2

In [266]: b = np.array([[1, 2, 3, 2, 1, 3],[2, 1, 3, 5, 4, 1]])

Since all values you want to change are greater than 2:

In [267]: b>2
Out[267]: 
array([[False, False,  True, False, False,  True],
       [False, False,  True,  True,  True, False]])
In [268]: np.where(b>2,0,b)
Out[268]: 
array([[1, 2, 0, 2, 1, 0],
       [2, 1, 0, 0, 0, 1]])

or pairing two tests:

In [271]: (b==1)|(b==2)
Out[271]: 
array([[ True,  True, False,  True,  True, False],
       [ True,  True, False, False, False,  True]])
In [272]: np.where((b==1)|(b==2),b,0)
Out[272]: 
array([[1, 2, 0, 2, 1, 0],
       [2, 1, 0, 0, 0, 1]])

We could change b in place, but I chose where so I can test various ideas without changing b .

isin uses a similar idea if one set is significantly smaller than the other, which is true in this case.