基于 2-D 布尔掩码屏蔽 3-D numpy 数组的行和列

Masking out row and column of 3-D numpy array based on 2-D boolean mask

对于如下所示的某些 3-D 立方体 numpy 数组:

import numpy as np
a = np.array([[[1,2,3],[4,5,6],[7,8,9]],[[10,11,12],[13,14,15],[16,17,18]],[[19,20,21],[22,23,24],[25,26,27]]])

array([[[ 1,  2,  3],
        [ 4,  5,  6],
        [ 7,  8,  9]],
       [[10, 11, 12],
        [13, 14, 15],
        [16, 17, 18]],
       [[19, 20, 21],
        [22, 23, 24],
        [25, 26, 27]]])

和一些二维布尔掩码数组,如下所示:

b = np.array([[0,1,1],[1,1,1],[1,1,0]])

array([[0, 1, 1],
       [1, 1, 1],
       [1, 1, 0]])

我想知道是否有一种方法可以使用 numpy 操作来计算结果,使得对于 b[i][j] = 0a[i,:,j] = 0a[i,j,:] = 0 的所有元素。保证 bn x n 并且 an x n x n。在上面的例子中,结果看起来像

array([[[ 0,  0,  0],
    [ 0,  5,  6],
    [ 0,  8,  9]],

   [[10, 11, 12],
    [13, 14, 15],
    [16, 17, 18]],

   [[19, 20,  0],
    [22, 23,  0],
    [ 0,  0,  0]]])
In [111]: b = np.array([[0,1,1],[1,1,1],[1,1,0]])

In [116]: I,J = np.nonzero(b==0)
In [117]: I,J
Out[117]: (array([0, 2]), array([0, 2]))

测试索引:

In [118]: a[I,:,J]
Out[118]: 
array([[ 1,  4,  7],
       [21, 24, 27]])
In [119]: a[I,J,:]
Out[119]: 
array([[ 1,  2,  3],
       [25, 26, 27]])

申请:

In [120]: a[I,:,J]=0
In [121]: a[I,J,:]=0
In [122]: a
Out[122]: 
array([[[ 0,  0,  0],
        [ 0,  5,  6],
        [ 0,  8,  9]],

       [[10, 11, 12],
        [13, 14, 15],
        [16, 17, 18]],

       [[19, 20,  0],
        [22, 23,  0],
        [ 0,  0,  0]]])