如何在 5x2x2 NumPy 数组中搜索特定的 2x2 矩阵？

Question

我有一个维度为 (5,2,2) 的 NumPy 数组，它是一系列 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])

然后，serie 中有一个用 2 填充的 2×2 矩阵。如何询问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
----------------------------------------