Как получить индексы максимальных элементов в массиве массивов NumPy?

Question

[Gfghhgg]

Имеется массив следующего вида:
(красным помечены наибольшие элементы, индексы которых надо получить)

То есть, для такого массива код должен возвращать что-то подобное:
(0, 2, 2), (1, 0, 1), (2, 1, 0), (3, 0, 0), (4, 0, 0)
Конечно это можно сделать циклом с np.argmax, но, возможно, кто-нибудь сможет найти более эффективное решение, реализованное только с помощью numpy (в моём случае очень важна скорость работы скрипта)
Возможно помогут эти функции: np.argmax, np.unravel_index

UPD:

Нашел способ достать сами элементы, но вот в чем проблема, мне нужно получить их индексы

Comments

[Gfghhgg]

Сергей Соколов, к сожалению нет, здесь нужна чуть более сложная комбинация numpy'ких функций

Answer 1

[Сергей Соколов]

мой массив

import numpy as np
x = np.random.random((5,3,3)) * 10 // 1

array([[[8., 8., 1.],
        [8., 1., 6.],
        [0., 7., 7.]],

       [[2., 7., 2.],
        [4., 5., 1.],
        [6., 8., 5.]],

       [[5., 1., 1.],
        [0., 1., 8.],
        [5., 7., 1.]],

       [[0., 5., 9.],
        [6., 7., 3.],
        [5., 0., 5.]],

       [[0., 1., 9.],
        [1., 2., 3.],
        [0., 4., 7.]]])

Раз максимум ищется среди 2-го уровня, можно изменить размерность, рассматривая каждую девятку как ряд:

y = x.reshape(5,9).argmax(axis=1)

# array([0, 7, 5, 2, 2])

Теперь эти индексы надо преобразовать обратно к паре координат:

f = lambda i: (i // 3, i % 3)
# или вернее так:
f = lambda i: (i // x.shape[1], i % x.shape[2])

(a, b) = f(y)

# array([0, 2, 1, 0, 0]), array([0, 1, 2, 2, 2])

Остаётся выдергивать из обоих массивов пары с одинаковыми индексами.

result = np.empty((a.size + b.size,), dtype=a.dtype)
result[0::2] = a
result[1::2] = b
result.reshape(5,2)

array([[0, 0],
       [2, 1],
       [1, 2],
       [0, 2],
       [0, 2]])

Comments

[Gfghhgg]

А вот это действительно интересное решение, спасибо большое!

[Gfghhgg]

Только вашу самописную функцию можно заменить на боолее производительную из numpy

Как видите вывод одинаков
Возьмите себе на заметку вдруг пригодится)
Ещё раз спасибо!

[Сергей Соколов]

Gfghhgg, вы сравнивали производительность, строили графики?
Или просто предположили, что реализация из numpy заведомо быстрее?
Просто нативное-питоновское быстрее нумпи – про lambda на SO люди заморочились сравнить этот и ещё 5 вариантов.

[Gfghhgg]

Так же мне удалось реализовать выдёргивания через комбинацию NumPy функций

[Gfghhgg]

И, вуаля, вот итоговое решение в одну строку

[Сергей Соколов]

Gfghhgg, ещё бы не картинками, а текстами, которые проиндексируются поисковиками и кому-то помогут потом.

[longclaps]

Что это, Бэримор?

f = lambda i: (i // x.shape[1], i % x.shape[2]) # вот-вот это что???

Ну и, конечно, это.

[Сергей Соколов]

longclaps, что-что, овсянка давно я питон не трогал, вот что. ТС нашёл более элегантную функцию numpy.unravel_index() для этого.

Answer 2

[kephirik]

Решение:
Пусть x некоторый 2мерный массив
np.where(x == np.max(x))