在numpy数组中进行设置的布尔索引

Question

我有形状为a的numpy数组m,n,3和形状为mask的{​​{1}}。尝试设置m,n时出现错误

a[mask > 0] = (255, 0, 0)

如何在numpy中实现此功能？

Answer 1

a = np.random.randint(0,2,size=(100,200,3))
mask = np.random.randint(0,2,size=(100,200))
a[mask > 0] = [255,0,0]
# Test
assert np.equal(a[mask > 0], 
    np.ones_like(a[mask > 0])*[255,0,0]).sum() == np.multiply(*a[mask > 0].shape)