重塑/扩展numpy蒙版

Question

I want to select elements in b given a condition based on a . However, the mask msk does not quite capture the output. Given the below,

>>> b = np.asarray([[[1, 11], [2, 12]], [[3, 13], [4, 14]], [[5, 15], [6, 16]], [[7, 17], [8, 18]]])
>>> b.shape
(4, 2, 2)
>>> b
array([[[ 1, 11],
        [ 2, 12]],

       [[ 3, 13],
        [ 4, 14]],

       [[ 5, 15],
        [ 6, 16]],

       [[ 7, 17],
        [ 8, 18]]])
>>> a = np.asarray([[0, 2], [3, 0], [0, 6], [0, 8]])
>>> a.shape
(4, 2)
>>> a
array([[0, 2],
       [3, 0],
       [0, 6],
       [0, 8]])
>>> msk = a > 0
>>> b[msk]
array([[ 2, 12],
       [ 3, 13],
       [ 6, 16],
       [ 8, 18]])
>>> msk
array([[False,  True],
       [ True, False],
       [False,  True],
       [False,  True]], dtype=bool)

I'd like to return [[11, 12], [3, 4], [15, 16], [17, 18]] instead of what b[msk] is currently returning. Intuitively, I think another dimension has to be added into msk converting it to the shape of b , ie (4, 2, 2) , but just calling msk.reshape didn't do the trick for me. I think with a mask like this, the problem would be solved. How can I convert the actual mask?

>>> msk
array([[[False,  True], [False,  True]],
       [[ True, False], [ True, False]],
       [[False,  True], [False,  True]],
       [[False,  True], [False,  True]]], dtype=bool)

Answer 1

Not sure if this is the best way, but:

>>> transposed = np.transpose(b, (2, 0, 1))
>>> np.hstack((transposed[0][:, None], transposed[1][:, None]))[msk]
array([[11, 12],
       [ 3,  4],
       [15, 16],
       [17, 18]])

Edit: With OP's suggestion we can simplify it to:

>>> np.transpose(b, (0, 2, 1))[msk]
array([[11, 12],
       [ 3,  4],
       [15, 16],
       [17, 18]])