CODEWITHSUNDEEP
pythonnumpy
user1101010
user1101010

Reputation: 187

Fast way to replace elements by zeros corresponding to zeros of another array

Suppose we have two numpy arrays

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

The goal is to set elements of b at the indices where a is 0. That is, we want to get an array

[ [1, 2, 0], [4, 0, 0], [7, 8, 0] ]

What is a fast way to achieve this?

I thought about generate a mask by $a$ first and then replace the values of b by this mask. But got lost on how to do this?

Upvotes: 3

Views: 601

Answers (2)

user3483203
user3483203

Reputation: 51175

You can use a masked_array here:

np.ma.masked_array(b, a==0).filled(0)

array([[1, 2, 0],
       [4, 0, 0],
       [7, 8, 0]])

Helpful if you don't want to modify b in place. You can replace filled(0) with whatever you'd like to set the elements equal to.

Upvotes: 1

wim
wim

Reputation: 363384

This is array assignment:

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

>>> a==0  # This is a boolean mask, True where the elements of `a` are zero
array([[False, False,  True],
       [False,  True,  True],
       [False, False,  True]])

>>> b[a==0] = 0  # So this is a masked assignment statement
>>> b
array([[1, 2, 0],
       [4, 0, 0],
       [7, 8, 0]])

Documented here.

Upvotes: 3

Related Questions