Vainmonde De Courtenay
Vainmonde De Courtenay

Reputation: 386

Replace numpy array with another array according to a mask

I have two numpy arrays A and B and a mask mask of booleans (True/False), all of identical dimensions. I want to replace the elements in A with those of B where the corresponding element of mask is True; where the corresponding element of mask is False I want to retain the original element of A. How can I do this?

Example:

# Input
A = np.arange(9).reshape(3,3)
B = A*10
mask = np.array([[True, True, False], [False, True, False], [False, False, True]])

# Output
desired_output = np.array([[0, 10, 2], [3, 40, 5], [6, 7, 80]])

Upvotes: 0

Views: 1395

Answers (2)

hpaulj
hpaulj

Reputation: 231738

Simply:

In [54]: A[mask]=B[mask]
In [55]: A
Out[55]: 
array([[ 0, 10,  2],
       [ 3, 40,  5],
       [ 6,  7, 80]])

np.putmask(A,mask,B) also works.

Upvotes: 3

Mario Abbruscato
Mario Abbruscato

Reputation: 819

Try map

def repl(a,b,m):
    return b if m else a

desired_output = list(map(repl,A,B,mask))

Upvotes: 0

Related Questions