NumPy Argmin and Argmax

numpy
Published

March 6, 2023

Understanding argmin()

The argmin() function returns the index of the minimum value along a specified axis of a NumPy array. If multiple minimum values exist, it returns the index of the first occurrence.

import numpy as np

arr = np.array([3, 1, 4, 1, 5, 9, 2, 6])

min_index = np.argmin(arr)
print(f"The index of the minimum value is: {min_index}")  # Output: 1

#argmin on a 2D array
arr_2d = np.array([[1, 5, 2], [8, 3, 9], [4, 7, 6]])
min_index_row = np.argmin(arr_2d, axis=0) #minimum index along each column
print(f"The indices of the minimum value along each column are: {min_index_row}") # Output: [0 1 0]

min_index_col = np.argmin(arr_2d, axis=1) #minimum index along each row
print(f"The indices of the minimum value along each row are: {min_index_col}") # Output: [0 1 0]

Understanding argmax()

Similarly, argmax() finds the index of the maximum value along a specified axis. Again, if multiple maximum values exist, it returns the index of the first one encountered.

import numpy as np

arr = np.array([3, 1, 4, 1, 5, 9, 2, 6])

max_index = np.argmax(arr)
print(f"The index of the maximum value is: {max_index}")  # Output: 5

#argmax on a 2D array
arr_2d = np.array([[1, 5, 2], [8, 3, 9], [4, 7, 6]])
max_index_row = np.argmax(arr_2d, axis=0) #maximum index along each column
print(f"The indices of the maximum value along each column are: {max_index_row}") # Output: [1 0 1]

max_index_col = np.argmax(arr_2d, axis=1) #maximum index along each row
print(f"The indices of the maximum value along each row are: {max_index_col}") # Output: [1 2 2]

Handling Multi-Dimensional Arrays

Both argmin() and argmax() gracefully handle multi-dimensional arrays. By specifying the axis argument, you can control whether the minimum/maximum is found along rows (axis=0), columns (axis=1), or other dimensions. The function will then return an array of indices, one for each row or column.

Beyond Simple Arrays: Practical Applications

The power of argmin() and argmax() extends beyond simple numerical arrays. They’re invaluable for tasks such as:

  • Image Processing: Identifying the location of the brightest or darkest pixel.
  • Machine Learning: Finding the class with the highest predicted probability.
  • Data Analysis: Locating extreme values in datasets.
  • Optimization: Determining the index of the best solution in a search space.

These functions are fundamental building blocks for numerous data manipulation and analysis tasks in Python, making them essential tools for any data scientist or programmer working with NumPy.