Как получить из трехмерного массива 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, 0, 2], точнее я бы хотел получить массив вида
[[False, False, True], [False, False, False]]

как это сделать без цикла?

P.S. - кстати в разделе Python это 28000 - ный вопрос)

Comments

[MinTnt]

Ну, а если бы это был обычный трехмерный массив, без numpy...

a = [[[1, 0, 0], [2, 0, 1], [2, 0, 2]], [[0, 0, 0], [2, 0, 3], [0, 0, 0]]]
k = [2, 0, 2]

print([*map(lambda x: [*map(lambda m: m==k,x)] ,a)]

Ну а так, решение ниже

Answer 1

[Vindicar]

(a == [2,0,2]).all(axis=2)
a == [2,0,2] даст тебе трёхмерный массив булевых значений - результат поэлементного сравнения, при этом сравнение будет идти по последнему измерению.
all(axis=2) сгруппирует этот массив по третьему измерению, при этом группа будет истинной, только если все значения в группе истинны. Т.е. только если в этой строке все значения совпали.

>>> a = np.array([[[1, 0, 0], [2, 0, 1], [2, 0, 2]], [[0, 0, 0], [2, 0, 3], [0, 0, 0]]])
>>> a == [2,0,2]
array([[[False,  True, False],
        [ True,  True, False],
        [ True,  True,  True]],

       [[False,  True, False],
        [ True,  True, False],
        [False,  True, False]]])
>>> (a == [2,0,2]).all(axis=2)
array([[False, False,  True],
       [False, False, False]])

Comments

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

Понял, спасибо