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

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

UPD:
5e7b43fcb902e378433006.png
Нашел способ достать сами элементы, но вот в чем проблема, мне нужно получить их индексы
  • Вопрос задан
  • 213 просмотров
Решения вопроса 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]])
Ответ написан
Пригласить эксперта
Ваш ответ на вопрос

Войдите, чтобы написать ответ

Войти через центр авторизации
Похожие вопросы