NumPy ndarray.argpartition()
The numpy.ndarray.argpartition() method returns the indices that would partition an array along a given axis.
It partially sorts the array such that the element at the kth position is in its final sorted position,
and all elements before it are smaller, while all elements after it are larger.
Syntax
ndarray.argpartition(kth, axis=-1, kind='introselect', order=None)Parameters
| Parameter | Type | Description |
|---|---|---|
kth | int or sequence of ints | Element index or indices to partition. It determines where the array should be partitioned. |
axis | int, optional | Axis along which to perform partitioning. Default is -1 (last axis). |
kind | {‘introselect’}, optional | Selection algorithm to use. Currently, only 'introselect' is available. |
order | list of str, optional | Used when partitioning structured arrays based on multiple fields. |
Return Value
Returns an array of indices indicating the partitioned ordering of elements. The elements at kth index are placed
in their correct final positions, while others are only partially ordered.
Examples
1. Basic Usage of ndarray.argpartition()
In this example, we apply argpartition() to a 1D array to find the index positions after partitioning around a specific element.
import numpy as np
# Create a 1D array
arr = np.array([30, 10, 20, 50, 40])
# Partition the array around the 2nd smallest element (index 2 in sorted order)
indices = arr.argpartition(2)
print("Original array:", arr)
print("Partitioned indices:", indices)
print("Rearranged array based on partitioning:", arr[indices])Output:
Original array: [30 10 20 50 40]
Partitioned indices: [1 2 0 4 3]
Rearranged array based on partitioning: [10 20 30 40 50]The argpartition(2) ensures that the element at index 2 (after sorting) is in its correct position,
while elements before it are smaller and elements after it are larger, but not necessarily sorted.
2. Using ndarray.argpartition() Along an Axis
Here, we apply argpartition() to a 2D array along a specific axis.
import numpy as np
# Create a 2D array
arr = np.array([[30, 10, 50],
[20, 40, 60]])
# Partition the array along axis 1 at index 1
indices = arr.argpartition(1, axis=1)
print("Original array:\n", arr)
print("Partitioned indices:\n", indices)
# Rearrange elements based on indices
print("Rearranged array:\n", np.take_along_axis(arr, indices, axis=1))Output:
Original array:
[[30 10 50]
[20 40 60]]
Partitioned indices:
[[1 0 2]
[0 1 2]]
Rearranged array:
[[10 30 50]
[20 40 60]]Partitioning ensures that the element at index 1 in each row is in its final sorted position, while values before it are smaller and values after it are larger.
3. Selecting Multiple Elements with ndarray.argpartition()
We can pass multiple indices to partition multiple elements correctly.
import numpy as np
# Create a 1D array
arr = np.array([100, 20, 80, 50, 40, 10])
# Partition the array around the 2nd and 4th smallest elements
indices = arr.argpartition([2, 4])
print("Original array:", arr)
print("Partitioned indices:", indices)
print("Rearranged array based on partitioning:", arr[indices])Output:
Original array: [100 20 80 50 40 10]
Partitioned indices: [5 1 4 3 2 0]
Rearranged array based on partitioning: [10 20 40 50 80 100]Here, the elements at index 2 and 4 (when fully sorted) are placed in their correct final positions, while others are only partially sorted.
4. Finding the k Smallest Elements Efficiently
argpartition() is useful for efficiently finding the k smallest elements in an array.
import numpy as np
# Create a 1D array
arr = np.array([100, 20, 80, 50, 40, 10])
# Get indices of the 3 smallest elements
k = 3
indices = arr.argpartition(k)[:k]
# Extract the k smallest elements
smallest_values = arr[indices]
print("Original array:", arr)
print("Indices of k smallest elements:", indices)
print("k smallest elements:", smallest_values)Output:
Original array: [100 20 80 50 40 10]
Indices of k smallest elements: [5 1 4]
k smallest elements: [10 20 40]The function efficiently finds the k smallest elements without fully sorting the array, improving performance.