用与另一个数组的零对应的零替换元素的快速方法

Question

假设我们有两个numpy数组

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的索引处设置b的元素。也就是说，我们要获得一个数组

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

实现此目标的快速方法是什么？

我考虑过先用$ a $生成一个掩码，然后用该掩码替换b的值。但是迷路了吗？

Answer 1

这是数组分配：

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

记录here。

Answer 2

您可以在此处使用masked_array：

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

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

如果您不想就地修改b，则很有用。您可以将filled(0)替换为您想要将元素设置为等于的任何内容。