compare two NumPy array and edit one of them-CodePudding

a = array([1,2,3,4,5,6])
b = array([0,1,1,0,1,0])

I would like to compare the arrays and then edit array a with the condition, if value in b is 1 the value in a should stay the same, but if value in b is 0, the value in a should be 0 too.

so a should look like this:

a = array([0,2,3,0,5,0])

At the moment I'm trying to do it with this code:

for x,y in np.nditer([a,b], op_flags=['readwrite']):
                if y != 1.0:
                    x[...] = 0.0

but it's very slow. Is there a faster solution?

Thank you all for your help!

CodePudding user response：

Hope the below approach helps:

import numpy as np
a = np.array([1,2,3,4,5,6])
b = np.array([0,1,1,0,1,0])
print(a*b)

Output:

[0 2 3 0 5 0]

CodePudding user response：

You can either multiply both arrays or use masking like this:

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

a[b == False] = 0 # Equivalent to a[b == 0] = 0