Чтобы получить индекс максимального элемента массива NumPy по определенной оси, вы можете использовать несколько методов. Вот некоторые часто используемые подходы:

1. Использование argmax(): функция argmax()из NumPy возвращает индексы максимальных значений вдоль заданной оси. Указав нужную ось, вы можете найти индекс максимального элемента по этой оси. Вот пример:

import numpy as np
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
max_index = np.argmax(arr, axis=0)
print(max_index)

Выход:

[2 2 2]

В этом примере argmax()возвращает индекс максимального значения по нулевой оси (по столбцу), который равен [2, 2, 2].

2. Использование amax()и where(): функция amax()возвращает максимальное значение вдоль указанной оси, а <Функция code>where()возвращает индексы, в которых выполняется заданное условие. Объединив эти две функции, можно найти индекс максимального элемента. Вот пример:

import numpy as np
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
max_value = np.amax(arr, axis=0)
max_indices = np.where(arr == max_value)
print(max_indices)

Выход:

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

В этом примере amax()возвращает максимальные значения по нулевой оси: [7, 8, 9]. Затем where()находит индексы, в которых встречаются эти максимальные значения, в результате чего получается (array([2, 2, 2]), array([0, 1, 2])).

3. Использование argpartition(). Функция argpartition()позволяет частично отсортировать массив вдоль указанной оси. Разделив массив таким образом, чтобы максимальный элемент находился в нужной позиции, вы можете получить индекс этого максимального элемента. Вот пример:

import numpy as np
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
max_indices = np.argpartition(arr, -1, axis=0)[-1]
print(max_indices)

Выход:

[2 2 2]

В этом примере argpartition()разбивает массив по нулевой оси так, чтобы максимальный элемент находился в последней позиции, а [-1]извлекает индексы этой позиции, в результате чего получается [2, 2, 2].