Deleting row in numpy array based on condition

Question

I have a 2D numpy array of shape [6,2] and I want to remove the subarrays with the third element containing 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]]) 

My code is positives = gt_boxes[np.where(gt_boxes[range(gt_boxes.shape[0]),2] != 0)]

It works but is there a simplified method to this?

Answer 1

You can use boolean indexing.

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

---

1. x[:, -1] will retrieve the last column

2. x[:, -1] != 0 returns a boolean mask

3. Use the mask to index into the original array