如何通过第一列值简单地过滤 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]]])

我怎样才能简单地从中提取一个 3D numpy 数组，其形状相同，轴 0 上有条件，例如选择轴 0 < 3 的那些“行”？ 一种天真的方式是

data[data[0]<3]

但这失败了：

IndexError：boolean 索引与维度 0 上的索引数组不匹配； 尺寸为 3，但对应的 boolean 尺寸为 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. ]]])