Home > Mobile >  How do I pass a list as changing condition in an array?
How do I pass a list as changing condition in an array?

Time:09-27

Let's say that I have an numpy array a = [1 2 3 4 5 6 7 8] and I want to change everything else but 1,2 and 3 to 0. With a list b = [1,2,3] a tried a[a not in b] = 0, but Python does not accept this. Currently I'm using a for loop like this:

c = a.unique()    

for i in c:
   if i not in b:
      a[a == i] = 0

Which works very slowly (Around 900 different values in a 3D array around the size of 1000x1000x1000) and doesn't fell like the optimal solution for numpy. Is there a more optimal way doing it in numpy?

CodePudding user response:

You can use numpy.isin() to create a boolean mask to use as an index:

np.isin(a, b)
# array([ True,  True,  True, False, False, False, False, False])

Use ~ to do the opposite:

~np.isin(a, b)
# array([False, False, False,  True,  True,  True,  True,  True])

Using this to index the original array lets you assign zero to the specific elements:

a = np.array([1,2,3,4,5,6,7,8])
b = np.array([1, 2, 3])

a[~np.isin(a, b)] = 0

print(a)
# [1 2 3 0 0 0 0 0]
  • Related