Как получить конкретный строки в 3-хмерном массиве NumPy?

Question

[Даниил Шевкунов]

У меня есть numpy массив:
a = np.array([[[1, 0, 0], [2, 0, 1], [2, 0, 2]], [[0, 0, 0], [2, 0, 3], [0, 0, 0]]])

Другой вид

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

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

мне надо получить все строки которые начинаются с 2 без цикла
Я думал сделать это как-то так:
a[: , a[: , :][0] == 2]
Но так возвращается это
([[2, 2, 2],
2, 0, 0]])

Что я делаю не так? Как правильно?

Answer 1

[dmshar]

Может вот так:

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]

Результат:

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