如果最后一个轴索引小于另一个2D数组中的值，则将3D numpy数组的值设置为0

Question

我有一个形状为（m，n，p）的3D数组a和a形状为（m，n）的2D数组idx 。 我希望最后一个轴索引小于idx相应元素的a中的所有元素都设置为0。

以下代码有效。 我的问题是：有没有更有效的方法？

a = np.array([[[1, 2, 3],
               [4, 5, 6]],

              [[7, 8, 9],
               [10, 11, 12]],

              [[21, 22, 23],
               [25, 26, 27]]])
idx = np.array([[2, 1],
                [0, 1],
                [1, 1]])
for (i, j), val in np.ndenumerate(idx):
    a[i, j, :val] = 0

结果是

array([[[ 0,  0,  3],
        [ 0,  5,  6]],

       [[ 7,  8,  9],
        [ 0, 11, 12]],

       [[ 0, 22, 23],
        [ 0, 26, 27]]])

Answer 1

使用broadcasting创建3D蒙版，然后使用boolean-indexing分配零-

mask = idx[...,None] > np.arange(a.shape[2])
a[mask] = 0

另外，我们也可以使用内置的NumPy进行更大范围的比较以获得该蒙版-

mask = np.greater.outer(idx, np.arange(a.shape[2]))

在给定的样本上运行-

In [34]: mask = idx[...,None] > np.arange(a.shape[2])

In [35]: a[mask] = 0

In [36]: a
Out[36]: 
array([[[ 0,  0,  3],
        [ 0,  5,  6]],

       [[ 7,  8,  9],
        [ 0, 11, 12]],

       [[ 0, 22, 23],
        [ 0, 26, 27]]])