Numpy, как найти индекс третьего измерения из трехмерного массива (например, х, у)

Question

Я использую NumPy.

Теперь у меня есть 3D-массив:

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

Я хочу найти в нем 2D координаты [4, 5, 6].

Ожидаемый результат: [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]))