根据条件删除numpy数组中的行

Question

我有一个形状为 [6,2] 的二维 numpy 数组，我想删除第三个元素包含 0 的子数组。

array([[0, 2, 1], #Input
       [0, 1, 1],
       [1, 1, 0],
       [1, 0, 2],
       [0, 2, 0],
       [2, 1, 2]])

array([[0, 2, 1], #Output
       [0, 1, 1],
       [1, 0, 2],
       [2, 1, 2]]) 

我的代码是positives = gt_boxes[np.where(gt_boxes[range(gt_boxes.shape[0]),2] != 0)]

它有效，但有简化的方法吗？

Answer 1

您可以使用布尔索引。

In [413]: x[x[:, -1] != 0]
Out[413]: 
array([[0, 2, 1],
       [0, 1, 1],
       [1, 0, 2],
       [2, 1, 2]])

---

1. x[:, -1] 将检索最后一列

2. x[:, -1] != 0 返回一个布尔掩码

3. 使用掩码索引到原始数组中