Numpy replace matrix elementwise with another matrix

Question

I have an n * x numpy matrix, which could look like this:

a = np.array([[1, 0], [0, 1]])

and I have another n * n numpy matrix, which look like this:

b = np.array([[2, 2], [2, 2]])

I would like to replace zero elements of a with the corresponding element of b so I would get:

[[1, 2],
 [2, 1]]

How can I do that?

Answer 1

You can just use a boolean mask:

mask = (a == 0)
a[mask] = b[mask]

This is efficient if you want to update the original array a since it only assigns to the zero elements, not the whole array.

Answer 2

You can use np.where:

np.where(a!=0, a, b)

array([[1, 2],
       [2, 1]])
​