Метод argmax

Метод argmax класса ndarray возвращает индекс максимального элемента в массиве. Если массив многомерный, можно указать ось axis, вдоль которой будет выполняться поиск. При наличии нескольких максимальных элементов метод возвращает индекс первого из них. Метод полезен при решении задач классификации, поиска экстремумов и анализа данных.

Синтаксис

arr.argmax([axis, out, keepdims])

Параметры метода:

axis - ось, вдоль которой выполняется поиск (по умолчанию массив сглаживается);
out - массив для вывода результата;
keepdims - если True, размерность результата сохраняется.

Пример

Давайте найдём индекс максимального элемента в одномерном массиве:

import numpy as np

arr = np.array([4, 2, 9, 1, 7])
res = arr.argmax()
print(res)

Результат выполнения кода:

2

Максимальный элемент 9 находится на индексе 2.

Пример

Рассмотрим поиск индекса максимального элемента в двумерном массиве без указания оси:

import numpy as np

arr = np.array([[1, 5, 3],
				[8, 2, 4]])
res = arr.argmax()
print(res)

Результат выполнения кода:

3

Массив сглаживается, и максимальный элемент 8 находится на позиции 3 в сглаженном представлении.

Пример

Найдём индексы максимальных элементов вдоль оси axis=0 (по столбцам):

import numpy as np

arr = np.array([[1, 5, 3],
				[8, 2, 4],
				[0, 6, 7]])
res = arr.argmax(axis=0)
print(res)

Результат выполнения кода:

[1 2 2]

Для каждого столбца возвращается индекс строки с максимальным элементом: в первом столбце максимальный элемент 8 в строке 1, во втором - 6 в строке 2, в третьем - 7 в строке 2.

Пример

Найдём индексы максимальных элементов вдоль оси axis=1 (по строкам):

import numpy as np

arr = np.array([[1, 5, 3],
				[8, 2, 4],
				[0, 6, 7]])
res = arr.argmax(axis=1)
print(res)

Результат выполнения кода:

[1 0 2]

Для каждой строки возвращается индекс столбца с максимальным элементом: в первой строке максимальный элемент 5 в столбце 1, во второй - 8 в столбце 0, в третьей - 7 в столбце 2.

Пример

Использование параметра keepdims для сохранения размерности:

import numpy as np

arr = np.array([[1, 5, 3],
				[8, 2, 4]])
res = arr.argmax(axis=0, keepdims=True)
print(res)
print(res.shape)

Результат выполнения кода:

[[1 0 1]]
(1, 3)

Результат сохраняет двумерную структуру, что полезно при дальнейших вычислениях с широковещанием.

Пример

Обработка случая с несколькими максимальными элементами (возвращается индекс первого):

import numpy as np

arr = np.array([3, 7, 2, 7, 1])
res = arr.argmax()
print(res)

Результат выполнения кода:

1

Максимальное значение 7 встречается на индексах 1 и 3, но метод возвращает индекс первого вхождения.

Смотрите также

метод argmin,
который возвращает индекс минимального элемента
метод max,
который возвращает значение максимального элемента
метод argsort,
который возвращает индексы отсортированных элементов
метод nonzero,
который возвращает индексы ненулевых элементов