如何在由掩码指定的较大数组的各个位置应用相同的一组值？

我有一个尺寸为Nx N的二维数组Y，以及一个K二进制掩码X的数组，每个尺寸都为M X M（所以X的形状为K x M x M）。“ X”中的每个二进制掩码都恰好有一个“ N x N”个1的补丁，其余元素为零。我想创建一个K x M x M数组Z，在X指定的位置使用Y`的元素，在其他所有位置将其设置为零。

例如，如果M = 3，N = 2，K = 3和

X = [[[1., 1., 0.],
     [1., 1., 0.],
     [0., 0., 0.]],

    [[0., 1., 1.],
     [0., 1., 1.],
     [0., 0., 0.]],

    [[0., 0., 0.],
     [0., 1., 1.],
     [0., 1., 1.]]]

和

Y = [[0.1, 0.2],
    [0.3, 0.4]]

那么Z应该是

Z = [[[0.1, 0.2, 0.],
     [0.3, 0.4, 0.],
     [0., 0., 0.]],

    [[0., 0.1, 0.2],
     [0., 0.3, 0.4],
     [0., 0., 0.]],

    [[0., 0., 0.],
     [0., 0.1, 0.2],
     [0., 0.3, 0.4]]]

我想使用Numpy或Pytorch中的函数来做到这一点。我认为使用numpy.where（）可以做到这一点，也许

Z = numpy.zeros((3,3,3))
Z[numpy.where(X == 1)] = Y

但这会导致形状不匹配错误。

什么是最简洁，最有效的方法（除了明确使用循环来设置值之外）？

问题来源：stackoverflow