Как искать конкретную матрицу 2x2 в массиве 5x2x2 NumPy?

Question

У меня есть массив NumPy размерности (5,2,2), который должен быть серией матриц 2 на 2. Как я могу спросить и имеет ли одна из этих матриц 2 на 2 конкретное значение?

Например, давайте возьмем ряд ниже

import numpy as np
zeros = np.zeros(4).reshape(2,2)
serie = np.array([zeros+1, zeros+2, zeros+3, zeros+4, zeros+5])

Затем, есть матрица 2 на 2, заполненная 2 с в serie. Как я могу спросить, содержит ли serie матрицу 2 на 2, заполненную 2, и получить ее индекс? В этом случае индекс будет равен 1, поскольку serie[1,:] является искомой матрицей.

Answer 1

Используя простое сравнение массивов, np.all и np.where:

import numpy as np

zeros = np.zeros(4).reshape(2,2)
serie = np.array([zeros+1, zeros+2, zeros+3, zeros+4, zeros+5])

to_find = zeros+2

index = np.where(np.all(serie == to_find, axis=(1, 2)))[0]
print(index)

Вывод:

[1]

Надежда это помогает!

----------------------------------------
System information
----------------------------------------
Platform:    Windows-10-10.0.16299-SP0
Python:      3.8.1
NumPy:       1.18.1
----------------------------------------