Python NumPy Sort, Unique, Search, and Filtering

Sorting, searching, and filtering are fundamental operations when working with numerical data in Python. NumPy provides a set of vectorized functions that make these operations fast and concise. This article focuses on python numpy sort unique search and filtering — the core functions for sorting arrays, extracting unique values, locating elements, and filtering based on conditions. You'll see how to apply these tools in real data workflows and understand the tradeoffs involved.

Sorting Arrays with numpy.sort and argsort

NumPy's sort function returns a sorted copy of an array without modifying the original. The function accepts an axis argument to sort along a specific dimension. For a 1D array, the default axis is -1, which sorts the entire array.

import numpy as np

arr = np.array([3, 1, 2, 5, 4])
sorted_arr = np.sort(arr)
print(sorted_arr)  # [1 2 3 4 5]
print(arr)         # [3 1 2 5 4] unchanged

When you need the indices that would sort the array, use argsort. This is useful when you have multiple arrays that share the same order, such as sorting one array and reordering another to match.

values = np.array([10, 5, 20])
indices = np.argsort(values)
print(indices)  # [1 0 2]
print(values[indices])  # [5 10 20]

For multi-dimensional arrays, sort sorts along the specified axis. Sorting along axis 0 sorts each column independently, while axis 1 sorts each row. The kind parameter lets you choose the sorting algorithm: 'quicksort' (default), 'mergesort', 'heapsort', and 'stable'. The stable option is useful when you need the original order of equal elements preserved.

Finding Unique Elements with numpy.unique

The unique function returns the sorted unique elements of an array. It also supports returning the indices of the first occurrences, the inverse indices, and the counts of each unique value.

arr = np.array([3, 1, 2, 3, 2, 1])
unique_vals, counts = np.unique(arr, return_counts=True)
print(unique_vals)  # [1 2 3]
print(counts)       # [2 2 2]

The return_inverse option gives an array that can be used to reconstruct the original array from the unique values. This is handy for label encoding or grouping.

unique_vals, inverse = np.unique(arr, return_inverse=True)
print(inverse)  # [2 0 1 2 1 0]
reconstructed = unique_vals[inverse]

For multi-dimensional arrays, unique flattens the input by default. Use the axis parameter to find unique rows or columns instead.

Searching Arrays: where, searchsorted, and argmax/argmin

numpy.where returns indices where a condition is true. It can also be used to select between two values based on a condition.

arr = np.array([1, 5, 3, 8, 2])
indices = np.where(arr > 3)
print(indices)  # (array([1, 3]),)

For sorted arrays, searchsorted is an efficient way to find insertion points that maintain order. It uses binary search, making it much faster than linear scanning for large arrays.

sorted_arr = np.array([1, 3, 5, 7])
positions = np.searchsorted(sorted_arr, [4, 6])
print(positions)  # [2 3]

argmax and argmin locate the index of the maximum or minimum value. These are often used to find the position of a peak or the closest match.

Filtering Arrays with Boolean Masks

Boolean indexing is the most direct way to filter NumPy arrays. You create a boolean mask by applying a comparison operator to the array, then use that mask to select elements.

arr = np.array([10, 20, 30, 40])
mask = arr > 20
filtered = arr[mask]
print(filtered)  # [30 40]

You can combine conditions with & (and), | (or), and ~ (not). Parentheses are required around each condition.

filtered = arr[(arr > 10) & (arr < 40)]
print(filtered)  # [20 30]

Boolean masks also work for assigning values. For example, you can replace all negative numbers with zero.

Combining Operations for Data Analysis

In practice, sorting, unique, search, and filtering are often used together. For instance, you might want to find the top three unique values in a dataset and then filter the original array to only those values.

data = np.array([7, 2, 7, 5, 2, 9, 9, 3])
unique_sorted = np.unique(data)
top_values = unique_sorted[-3:]  # largest three unique values
filtered = data[np.isin(data, top_values)]
print(filtered)  # [7 7 9 9]

np.isin is a convenient way to test membership and is often used with filtering. It returns a boolean mask indicating whether each element appears in a given set.

Performance and Memory Considerations

NumPy's vectorized operations are implemented in C and avoid Python-level loops, making them much faster for large arrays. However, some functions create intermediate arrays. For example, np.unique sorts the array internally, which has O(n log n) time complexity and requires additional memory for the sorted copy. When memory is tight, consider whether you need the full unique set or just a subset.

searchsorted is O(log n) and is ideal for repeated queries on a sorted array. If you need to search many values, using searchsorted on a pre-sorted array is far more efficient than a linear scan with np.where.

Boolean masks create a new array of the same shape as the original, so filtering large arrays can double memory usage. If memory is a constraint, consider using np.compress with a precomputed mask, though it behaves similarly.

Common Pitfalls and Edge Cases

One common mistake is assuming that np.sort modifies the array in place. It does not; use arr.sort() for in-place sorting. Another is using np.unique on a multi-dimensional array without specifying axis, which flattens the data unexpectedly.

When combining conditions with & and |, forgetting parentheses leads to a ValueError because Python's operator precedence treats them incorrectly. Always wrap each comparison in parentheses.

searchsorted assumes the array is sorted. If you pass an unsorted array, the results are meaningless. Ensure the array is sorted before calling it.

For floating-point arrays, np.unique uses exact equality. Values that are extremely close but not identical will be treated as distinct. If you need tolerance-based uniqueness, you may need to round or use a custom approach.