numpy如何从三维数组（例如x，y）中找到三维索引

Question

我正在使用numpy。

现在我有一个3D阵列：

ary = array([[[1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]]])

我想在其中找到[4, 5, 6]的2D坐标。

预期结果是：[0,1], [0,4]。

但是如果我使用where，结果将是：

np.where(s==[4,5,6])
(array([0, 0, 0, 0, 0, 0]), array([1, 1, 1, 4, 4, 4]), array([0, 1, 2, 0, 1, 2]))

那么，我该怎么做才能得到像[0,1], [0,4]这样的结果？

Answer 1

np.where((ary == [4,5,6]).all(axis=-1))
# (array([0, 0]), array([1, 4]))