根据掩码将 numpy 数组替换为另一个数组

Question

我有两个 numpy 数组 A 和 B 以及一个布尔值掩码 mask (True/False)，所有这些数组的维度相同。我想将A中的元素替换为B中的元素，其中mask对应的元素是True；其中mask对应的元素是False 我想保留A原来的元素。我该怎么做？

示例：

# 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]])

Answer 1

简单地说：

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) 也有效。

Answer 2

尝试map

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

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

根据掩码将 numpy 数组替换为另一个数组

Replace numpy array with another array according to a mask

python

arrays

numpy

mask