Revision history [back]

The argmax operation can be performed quickly in numpy using the argmax function. The argmax function returns the indices of the maximum values along an axis.

Here is an example:

import numpy as np a = np.array([[5, 2, 3], [7, 1, 4], [0, 8, 6]]) # Find the index of the maximum value along axis 0 max_idx = np.argmax(a, axis=0) print(max_idx) # Output: [1 2 2] # Find the index of the maximum value along axis 1 max_idx = np.argmax(a, axis=1) print(max_idx) # Output: [0 0 1] 

In this example, we have a 3x3 numpy array a and we are finding the index of the maximum value along axis 0 and axis 1 using the argmax function. We get the output [1 2 2] when we find the maximum value along axis 0, which means that the maximum values are located at index 1 in the first column, index 2 in the second column, and index 2 in the third column. Similarly, we get the output [0 0 1] when we find the maximum value along axis 1, which means that the maximum values are located at index 0 in the first row, index 0 in the second row, and index 1 in the third row.