Partial sorting with np.partition() in NumPy

Introduction

Partial sorting helps you quickly find some smallest or largest values without sorting everything. It saves time and effort.

Finding the top 3 lowest temperatures from a large list of daily temperatures.

Getting the 5 fastest runners' times from a race without sorting all participants.

Selecting the smallest 10 prices from a big list of product prices.

Finding the median value in a large dataset efficiently.

Syntax

NumPy

np.partition(array, kth, axis=-1, kind='introselect', order=None)

array is your data to partially sort.

kth is the index or indices of the elements to place correctly.

Examples

This puts the element that would be at index 2 in a sorted array into its correct place. Elements before it are smaller or equal, after it are larger or equal.

NumPy

import numpy as np
arr = np.array([7, 2, 5, 3, 9])
np.partition(arr, 2)

This places elements at indices 1 and 3 in their sorted positions, partially sorting the array around these points.

NumPy

np.partition(arr, [1, 3])

After partitioning, the first 3 elements are the smallest 3 values but not fully sorted.

NumPy

np.partition(arr, 2)[0:3]

Sample Program

This code finds the 4 smallest numbers in the array quickly using np.partition. The array is rearranged so that the element at index 3 is in the correct sorted position. Elements before it are smaller or equal, but not fully sorted.

NumPy

import numpy as np

# Create an array of random numbers
data = np.array([12, 3, 5, 7, 19, 1, 8])

# Partially sort to find the 4 smallest elements
part_sorted = np.partition(data, 3)

# The first 4 elements now contain the 4 smallest values (indices 0 to 3)
smallest_four = part_sorted[:4]

print('Original array:', data)
print('Partially sorted array:', part_sorted)
print('Smallest four elements:', smallest_four)

Important Notes

Note: The elements before the kth index are not fully sorted, just guaranteed to be smaller or equal.

You can use np.partition to speed up finding smallest or largest values without full sorting.

Summary

Partial sorting quickly finds some smallest or largest values.

np.partition rearranges the array so the kth element is in the right place.

Elements before and after are only partially ordered, not fully sorted.