numpy.ndarray.argmax

ndarray.argmax(axis=None, out=None)

Return indices of the maximum values along the given axis of a.

Parameters:

axis : int, optional

Axis along which to operate. By default flattened input is used.

out : ndarray, optional

Alternative output array in which to place the result. Must be of the same shape and buffer length as the expected output.

Returns:

index_array : ndarray

An array of indices or single index value, or a reference to out if it was specified.

Examples

>>> a = np.arange(6).reshape(2,3)
>>> a.argmax()
5
>>> a.argmax(0)
array([1, 1, 1])
>>> a.argmax(1)
array([2, 2])