Python NumPy: How to Find the K Smallest Values in a NumPy Array in Python

When analyzing numerical data, a common task is finding the k smallest values in an array - whether it's identifying the lowest temperatures, cheapest products, or smallest distances. While you could sort the entire array and slice it, NumPy and Python offer more efficient approaches that avoid the overhead of a full sort.

This guide covers four methods for finding the k smallest elements, ranging from the most performant to the most straightforward.

Using np.partition() (Recommended)

np.partition() is the fastest method for finding k smallest values because it performs a partial sort. It rearranges the array so that the element at the k-th position is where it would be in a fully sorted array, all smaller elements are moved before it, and all larger elements are moved after it - without sorting either group internally.

import numpy as np

arr = np.array([23, 12, 1, 3, 4, 5, 6])
k = 4

result = np.partition(arr, k)[:k]

print("Array:", arr)
print(f"{k} smallest elements:", result)

Output:

Array: [23 12  1  3  4  5  6]
4 smallest elements: [1 3 4 5]

Why np.partition() Is Faster

For large arrays where you only need a small number of elements, the difference can be significant.

Using np.argpartition() for Indices

np.argpartition() works like np.partition(), but instead of returning the values, it returns the indices of the k smallest elements. This is useful when you need to know where the smallest values are located in the original array.

import numpy as np

arr = np.array([23, 12, 1, 3, 4, 5, 6])
k = 4

indices = np.argpartition(arr, k)[:k]
values = arr[indices]

print("Array:", arr)
print("Indices of k smallest:", indices)
print(f"{k} smallest values:", values)

Output:

Array: [23 12  1  3  4  5  6]
Indices of k smallest: [2 3 4 5]
4 smallest values: [1 3 4 5]

Getting Sorted Results with Indices

If you need both the indices and values in sorted order:

import numpy as np

arr = np.array([23, 12, 1, 3, 4, 5, 6])
k = 4

indices = np.argpartition(arr, k)[:k]

sorted_order = indices[np.argsort(arr[indices])]

print("Sorted indices:", sorted_order)
print("Sorted values:", arr[sorted_order])

Output:

Sorted indices: [2 3 4 5]
Sorted values: [1 3 4 5]

Using np.sort() (Simple but Slower)

The most straightforward approach is to fully sort the array and then slice the first k elements. While less efficient than np.partition(), it's easy to read and guarantees the result is in sorted order.

import numpy as np

arr = np.array([23, 12, 1, 3, 4, 5, 6])
k = 4

result = np.sort(arr)[:k]

print("Array:", arr)
print(f"{k} smallest elements (sorted):", result)

Output:

Array: [23 12  1  3  4  5  6]
4 smallest elements (sorted): [1 3 4 5]

Using heapq.nsmallest() (Python Standard Library)

Python's built-in heapq module provides nsmallest(), which uses a heap data structure to efficiently find the k smallest elements. This works on any iterable and doesn't require NumPy.

import numpy as np
import heapq

arr = np.array([23, 12, 1, 3, 4, 5, 6])
k = 4

result = heapq.nsmallest(k, arr.tolist())

print("Array:", arr)
print(f"{k} smallest elements:", result)

Output:

Array: [23 12  1  3  4  5  6]
4 smallest elements: [1, 3, 4, 5]

Practical Example: Finding K Smallest in a 2D Array

In real-world scenarios, you might need the k smallest values from a multi-dimensional array. Here's how to handle that:

import numpy as np

data = np.array([
    [45, 12, 67],
    [3, 89, 23],
    [56, 7, 41]
])
k = 4

flat = data.ravel()
result = np.sort(np.partition(flat, k)[:k])

print("2D Array:")
print(data)
print(f"\n{k} smallest values across entire array:", result)

Output:

2D Array:
[[45 12 67]
 [ 3 89 23]
 [56  7 41]]

4 smallest values across entire array: [ 3  7 12 23]

Comparison of Approaches

Finding the K Largest Values

All these methods work for finding the k largest values too - simply adjust the approach:

import numpy as np

arr = np.array([23, 12, 1, 3, 4, 5, 6])
k = 3

result = np.partition(arr, -k)[-k:]

print(f"{k} largest elements:", np.sort(result))

Output:

3 largest elements: [ 6 12 23]

Using -k as the partition index places the k largest elements at the end of the partially sorted array.

Conclusion

For finding the k smallest values in a NumPy array, np.partition() is the best choice for most scenarios - it runs in linear time and scales efficiently to large datasets.

• Use np.argpartition() when you also need the indices of those elements.
• Fall back to np.sort() for simplicity on small arrays.
• Use heapq.nsmallest() when working outside the NumPy ecosystem.

Remember that np.partition() doesn't guarantee sorted order within the k elements, so add a secondary sort on the result if ordering matters.