Put zeros in the given index values of a 2D array in Python

Question

I have arrays

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

B = np.array([
  [3,2,1],
  [9,2,8],
  [2,1,8]
])

Doing a = np.argwhere(A > 4) gives me an array of position/indexes of values in array A that are greater than 4 i.e. [[0 1], [0 2], [2 0]].

I need to use these indexes/position from a = np.argwhere(A > 4) to replace the values in array B to zero at these specific position i.e. array B should now be

B = np.array([
  [3,0,0],
  [9,2,8],
  [0,1,8]
])

I am big time stuck any help with this will be really appreciated.

Thank You :)

Answer 1

In general, though, note that the indices returned by np.where are meant to be applied to numpy.ndarray objects, so you could have done:

B[np.where(A > 4)] = 0

Generally I don't use np.where with a condition like this, I just use the boolean mask directly, as in John Zwinck's answer. But it is probably important to understand that you could

>>> B[np.where(A > 4)] = 0
>>> B
array([[3, 0, 0],
       [9, 2, 8],
       [0, 1, 8]])

Answer 2

It should be as simple as:

B[A > 4] = 0