NumPy partition: function step-by-step examples


numpy partition step by step examples

In this blog post, we will delve into the intricacies of the 'numpy.partition()' function, exploring its syntax, applications, and providing step-by-step examples to deepen your understanding. Numpy is the One such function that stands out is 'numpy.partition()'.


Understanding 'numpy.partition()':

The 'numpy.partition()' function is designed to rearrange the elements of an array in such a way that the k-th element is in its sorted position within the array. This function is particularly useful when you only need a specific subset of elements in sorted order, without the need to sort the entire array.


Let's begin by exploring the key features of the NumPy partition() function and its

numpy.partition(arr, kth, axis, kind, order)

numpy.partition(arr, kth, axis=-1, kind='introselect', order=None)

  • 'arr': The input array to be partitioned.
  • 'kth': The index of the element that will be in its sorted position after the partition.
  • 'axis': The axis along which the partition is performed. The default is -1 (last axis).
  • 'kind': The algorithm used to perform the partition. The default is 'introselect', a hybrid algorithm.
  • 'order': If the array is structured (e.g., a structured dtype), this parameter specifies the order of the sort.

Now, let's dive into some step-by-step examples to illustrate the usage of 'numpy.partition()'.

Read also:


Harnessing the Power of NumPy partition() for Efficient Data Sorting

Sorting in Ascending Order

One of the primary use cases of the NumPy partition() function is sorting an array in ascending order and extracting the smallest k elements. This can be achieved with a concise and powerful one-liner:

import numpy as np

arr = np.array([4, 2, 9, 7, 5])
result = np.partition(arr, 3)

In this example, the array [4, 2, 9, 7, 5] is partitioned around the third smallest element, yielding the sorted result [2, 4, 5, 7, 9].

Detailed Explanation: Numpy Partition Function

Step 1: Original Array

arr = np.array([4, 2, 9, 7, 5])

Step 2: Calling numpy.partition()

result = np.partition(arr, 3)

Certainly! Let's visualize step by step how the numpy.partition() works with the given example.

Step 1: Original Array

arr = np.array([4, 2, 9, 7, 5])
The original array is [4, 2, 9, 7, 5].

Step 2: Calling numpy.partition()

result = np.partition(arr, 3)

The numpy.partition() function is called with 'kth=3', meaning we want the 3rd element (index 2) to be in its final sorted position after the partition.

Step 3: Partitioning the Array

Now, let's visualize how the partitioning happens:

The pivot element is chosen based on the kth parameter, which is the 3rd element (index 2) in this case.

Elements less than the pivot are moved to the left, and elements greater than the pivot are moved to the right.

Name Partitioned Array
Initial Configuration [4, 2, 9, 7, 5]
Pivot Element (3rd element) [4, 2, (9), 7, 5]
Partitioned Array [2, 4, (5), 7, 9]

Step 4: Result

The result of the 'numpy.partition()' operation is '[2, 4, 5, 7, 9]'. The array is now partially sorted, with the 3rd element (index 2) in its sorted position.

The elements to the left of the pivot (2 and 4) are smaller, and the elements to the right (7 and 9) are larger. The element at the pivot position (5) is in its final sorted position after the partition.



Sorting in Descending Order

To sort an array in descending order and extract the largest k elements, a slight modification is required:

import numpy as np

arr = np.array([4, 2, 9, 7, 5])
result_descending = np.partition(arr, -3)[::-1]



[9 7 5 4 2]

Here, we utilize array slicing '([::-1])' to reverse the order of the first three elements, achieving the desired descending order.

Example 1: Basic Usage

import numpy as np

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

# Partition the array around the 3rd element (kth=2)
result = np.partition(arr, 2)

print("Original array:", arr)
print("Partitioned array:", result)


Original array: [3 1 4 1 5 9 2 6 5 3 5]
Partitioned array: [1 1 2 3 3 4 5 6 5 9 5]

In this example, the third element (index 2) is in its sorted position, and the elements less than 2 are on its left, while elements greater than 2 are on its right.

Example 2: Multi-dimensional Arrays

import numpy as np

# Create a 2D array
arr_2d = np.array([[3, 1, 4], [1, 5, 9], [2, 6, 5]])

# Partition along the second axis (axis=1) around the 2nd element (kth=1)
result_2d = np.partition(arr_2d, 1, axis=1)

print("Original 2D array:")
print("\nPartitioned 2D array:")


Original 2D array:
[[3 1 4]
 [1 5 9]
 [2 6 5]]

Partitioned 2D array:
[[1 3 4]
 [1 5 9]
 [2 5 6]]

In this case, the elements along the second axis (columns) are partitioned independently.

Example 3: Using a Structured Array

import numpy as np

# Create a structured array with dtype
dtype = [('name', 'S10'), ('age', int), ('score', float)]
arr_structured = np.array([('chandra', 25, 90.5), ('kumar', 30, 85.0), ('ravi', 22, 95.5)], dtype=dtype)

# Partition the structured array based on the 'age' field
result_structured = np.partition(arr_structured, 1, order='age')

print("Original structured array:")
print("\nPartitioned structured array:")


Original structured array:
[(b'chandra', 25, 90.5) (b'kumar', 30, 85. ) (b'ravi', 22, 95.5)]

Partitioned structured array:
[(b'ravi', 22, 95.5) (b'chandra', 25, 90.5) (b'kumar', 30, 85. )]

In this example, the structured array is partitioned based on the 'age' field.


The 'numpy.partition()' function is a powerful tool for efficiently obtaining a partially sorted version of an array. By understanding its syntax and exploring examples, you can leverage this function to streamline your data manipulation tasks in a variety of scenarios. Whether you are working with one-dimensional or multi-dimensional arrays, or even structured data, 'numpy.partition()' provides a flexible and efficient solution.

Previous Post Next Post