numpy.argmax

numpy.argmax(a, axis=None, out=None)

Функция argmax() возвращает индекс максимального значения вдоль указанной оси.

Параметры:
a - массив NumPy или подобный массиву объект.
Исходный массив.
axis - целое число (необязательный параметр).
Определяет ось вдоль которой выполняется поиск элементов. Если равен None (по умолчанию), то поиск выполняется по сжатому до одной оси представлению исходного массива.
out - массив NumPy (необязательный параметр).
Позволяет сразу напрямую поместить результат в указанный массив, при условии, что он имеет подходящую форму и тип данных.
Возвращает:
ndarray - массив NumPy
массив индексов максимальных элементов в исходном массиве a с той же формой но удаленным размером по указанной в axis оси.
Смотрите так же:
nanargmax, argmin, nanargmin, unravel_index, take_along_axis

Замечание

В случаях, когда максимальному значению соответствует несколько элементов, то возвращается индекс элемента с первым вхождением.

Данная функция обладает эквивалентным методом класса ndarray, т.е. np.argmax(a) равносильно вызову метода a.argmax():

>>> a = np.random.randint(0, 20, 10)
>>> a
array([18,  4,  1, 19,  0, 14,  1, 17,  2,  6])
>>> 
>>> np.argmax(a)
3
>>> 
>>> a.argmax()
3


Примеры

По умолчанию axis = None, что соответствует поиску максимальных значений в плоском представлении массива:

>>> import numpy as np
>>> 
>>> a = np.random.randint(0, 20, size = (3, 5))
>>> a
array([[14,  7, 14, 15, 16],
       [ 6, 11,  8, 12, 15],
       [16, 19,  2,  7, 17]])
>>> 
>>> 
>>> np.argmax(a)
11
>>> 
>>> a.flat[11]
19

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

>>> np.argmax(a, axis = 0)
array([2, 2, 0, 0, 2], dtype=int32)

Получить сами элементы по данным индексам можно с помощью функции np.take_along_axis():

>>> ind_col = np.argmax(a, axis = 0)
>>> ind_col = np.expand_dims(ind_col, axis = 0)
>>> ind_col
array([[2, 2, 0, 0, 2]], dtype=int32)
>>> 
>>> 
>>> np.take_along_axis(a, ind_col, axis = 0)
array([[16, 19, 14, 15, 17]])

В том случае если мы ищем максимальный элемент на основе его индекса в плоском представлении массива, мы можем воспользоваться функцией np.unravel_index():

>>> ind = np.unravel_index(np.argmax(a), a.shape)
>>> ind
(2, 1)
>>> 
>>> a[ind]
19

Однако, возвращаемые индексы, как правило, используются для поиска соответствующих элементов в других массивах:

>>> a
array([[14,  7, 14, 15, 16],
       [ 6, 11,  8, 12, 15],
       [16, 19,  2,  7, 17]])
>>> 
>>> b = np.tile(np.arange(100, 103).reshape(3, 1), 5)
>>> b
array([[100, 100, 100, 100, 100],
       [101, 101, 101, 101, 101],
       [102, 102, 102, 102, 102]])
>>> 
>>> 
>>> np.take_along_axis(a, ind_col, axis = 0)
array([[16, 19, 14, 15, 17]])
>>> 
>>> np.take_along_axis(b, ind_col, axis = 0)
array([[102, 102, 100, 100, 102]])
>>> 
>>> 
>>> b[ind]
102