How to get specific rows in a 3 dimensional NumPy array?

Question

Python

I have a numpy array:
a = np.array([, ])

Another kind

([[[1, 0, 0],
[2, 0, 1],
[2, 0, 2]],

[[0, 0, 0],
[2, 0, 3],
[0, 0, 0]]])

i need to get all rows that start with 2 without a loop
i thought to do it something like this:
a[: , a[: , :][0] == 2]
But this returns this
([[2, 2, 2] ,
2, 0, 0]])

What am I doing wrong? How right?

Answer 1

Maybe like this:

import numpy as np
a = np.array([[[1, 0, 0], [2, 0, 1], [2, 0, 2]], [[0, 0, 0], [2, 0, 3], [0, 0, 0]]])
a[a[:,:,0]==2]

Result:

array([[2, 0, 1],
       [2, 0, 2],
       [2, 0, 3]])