python - 如何通过第一列值简单地过滤 3d numpy 数组？

Question

假设我有一个像这样的 3D numpy 数组:

 data = np.array([[[1,2,3,4],[1,2.5,3,5]],
                 [[116,230,450,430],[80,100,300,320]],
                 [[60,100,120,80],[50,80,100,90]]])

如何简单地从中提取具有相同形状且轴 0 上有条件的 3D numpy 数组，例如选择轴 0 < 3 的那些“行”？一个天真的方法是

data[data[0]<3]

但这失败了:

IndexError: bool 索引与维度 0 上的索引数组不匹配；维度为 3，但对应的 bool 维度为 2

Answer 1

请参阅上面的评论，但从您的数据来看，我猜测您想要任何值小于 3 的行。如果是这样，您可以这样做:

data[(data<3).any(axis=2)]
>>> array([[1. , 2. , 3. , 4. ],
           [1. , 2.5, 3. , 5. ]])

编辑1:

可以使用转置来匹配轴尺寸来实现解决方案:

data.T[(data[0]<3).any(axis=0).T].T
>>> array([[[  1. ,   2. ],
            [  1. ,   2.5]],

           [[116. , 230. ],
            [ 80. , 100. ]],

           [[ 60. , 100. ],
            [ 50. ,  80. ]]])

编辑2:

另一种不涉及转置的方法。敷面膜(data[0]<3).any(axis=0)到原始数据数组上，轴形状必须匹配。口罩形状为(4,)和data.shape = (3, 2, 4) ，所以我们需要将掩码应用到最后一个轴:

data[..., (data[0]<3).any(axis=0)]
>>> array([[[  1. ,   2. ],
            [  1. ,   2.5]],

           [[116. , 230. ],
            [ 80. , 100. ]],

           [[ 60. , 100. ],
            [ 50. ,  80. ]]])