Comparing 2 numpy arrays

Question

I have 2 numpy arrays and I want whenever element B is 1, the element in A is equal to 0. Both arrays are always in the same dimension:

A = [1, 2, 3, 4, 5]
B = [0, 0, 0, 1, 0]

I tried to do numpy slicing but I still can't get it to work.

B[A==1]=0

How can I achieve this in numpy without doing the conventional loop ?

Answer 1

First, you need them to be numpy arrays and not lists. Then, you just inverted B and A.

import numpy as np
A = np.array([1, 2, 3, 4, 5])
B = np.array([0, 0, 0, 1, 0])
A[B==1]=0 ## array([1, 2, 3, 0, 5])

If you use lists instead, here is what you get

A = [1, 2, 3, 4, 5]
B = [0, 0, 0, 1, 0]
A[B==1]=0 ## [0, 2, 3, 4, 5]

That's because B == 1 is False or 0 (instead of an array). So you essentially write A[0] = 0

Answer 2

Isn't it that what you want to do ?

A[B==1] = 0
A
array([1, 2, 3, 0, 5])