如何通過第一列值簡單地過濾 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. ]]])