Python numpy.argmax() function – Variants to know!

NumPy Argmax() Function

Hello, readers! In this article, we will be focusing on Python numpy.argmax() function in detail.

So, let us get started!! 🙂


Working on numpy.argmax() function

Python offers us a powerful tool in the form of a module to work on complex mathematical computations. Yes, you guessed it right 🙂

NumPy module, also known for its great complexity features offers us various data structures to deal with the data. Prominently we deal with the data in terms of Arrays and Matrices. Be it any computation problem or a real-life problem, analytics and preprocessing seem to be incomplete without NumPy as the mathematical model solution.

Today, we will be particularly discussing about NumPy argmax() function. This function helps us get the maximum elements at ease from within a matrix in the set of rows and columns. Here, it not only enables us to find the maximum element from the matrix, but also gives us the flexibility to choose the axis.

By axis, I mean to say the row or the column along which we want the maximum value to be fetched from the entire matrix.

Having understood this, in the upcoming section, will be focusing on the different variants offered within the numpy.argmax() function.


1. Fetch maximum element from a matrix using argmax()

On a basic scenario, the argmax() function enables us to have the maximum element from the entire array. But, when we do not specify any axis i.e. either by row or column, then it simply returns the count of elements at output.

For the same, if we wish to have the maximum element from the entire array without having to specify the axis, then will need to make use of the unravel_index() function as shown below.

Example:

At first, we need to create a NumPy array. Here, we have made use of arange() function with reshape() function to create a multi dimensional array of random elements within the specified range.

import numpy as np
arr = np.arange(9).reshape(3,3) 
print(arr)
print("Maximum element's index from the entire matrix")
pos = np.unravel_index(np.argmax(arr), arr.shape)
print(pos)
print(arr[pos])

Output:

[[0 1 2]
 [3 4 5]
 [6 7 8]]
Maximum element's index from the entire matrix
(2, 2)
8

2. Fetch maximum element for a specific row

Apart from fetching maximum element value/index from the entire array, while processing the data in terms of rows and columns, we may come across need to manipulate the data and get the maximum value across rows.

At this point, we need to add axis=1 as a parameter to the argmax() function as shown.

This function returns us the index of the maximum element across every row as output.

Example:

import numpy as np
arr = np.arange(9).reshape(3,3) 
print(arr)
print("Maximum element's index per row")
print(np.argmax(arr, axis=1))

Output:

[[0 1 2] 
 [3 4 5] 
 [6 7 8]]
Maximum element's index per row
[2 2 2]

3. Fetch maximum element for specific column

In order to get the maximum element across a column, we need to replace the value for axis with 0 i.e. axis=0.

With this, the argmax() function would return the index values of all the maximum elements present in each column of the matrix/array.

Example:

As seen below, here we have created a 3×3 array and then applied axis=0 for it to get the index of maximum element from every column as output.

import numpy as np
arr = np.arange(9).reshape(3,3) 
print(arr)
print("Maximum element's index per column")
print(np.argmax(arr, axis=0))

Output:

As elements 6, 7 and 8 are the maximum element per row, that is why we have the index values as [2, 2, 2].

[[0 1 2]
 [3 4 5]
 [6 7 8]]
Maximum element's index per column
[2 2 2]

Conclusion

By this, we have come to the end of this topic. Feel free to comment below, in case you come across any questions. For more such posts related to Python programming, Stay tuned with us. Till then, Happy Learning!! 🙂