重塑/扩展numpy面具

时间:2015-03-13 07:27:28

标签: python numpy

我想根据b选择a中的元素。但是,掩码msk并不能完全捕获输出。鉴于以下情况,

>>> 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)

我想返回[[11, 12], [3, 4], [15, 16], [17, 18]]而不是b[msk]当前正在返回的内容。直觉上,我认为必须将另一个维度添加到msk中,将其转换为b的形状,即(4, 2, 2),但只是调用msk.reshape没有做到骗局给我。我想戴着这样的面具,问题就解决了。如何转换实际面具?

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

1 个答案:

答案 0 :(得分:2)

不确定这是否是最佳方式,但是:

>>> 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]])

编辑:根据OP的建议,我们可以将其简化为:

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