在numpy中的3D數組中查找2D數組的索引

Question

我想在3d數組中找到2d數組。 為了在2d數組中找到1d數組，我可以使用np.where(np.all(a==b, axis=1))[0][0] 。

>>> import numpy as np
>>>
>>> a = np.array([[[1, 0, 0],
                   [0, 0, 0],
                   [0, 0, 0]],

                  [[0, 0, 0],
                   [0, 0, 0],
                   [0, 0, 0]]])
>>>
>>> b = np.array([[1, 0, 0],
                  [0, 0, 0],
                  [0, 0, 0]])
>>>
>>> a.find(b)
0

Answer 1

axis關鍵字接受元組，因此您可以簡單地執行以下操作：

np.where(np.all(a==b, axis=(1, 2)))[0][0]