Introduction
Understanding the intricacies of NumPy’s ndarray.argsort() method is essential for anyone dealing with numerical data in Python. This guide will take you through the basics to more advanced usage of the argsort() method with four illustrative examples. Whether you are sorting basic one-dimensional arrays, or you’re dealing with more complex multidimensional data, mastering argsort() is a crucial step in handling numerical datasets efficiently.
What does ndarray.argsort() return?
The argsort() method in NumPy returns the indexes that would sort an array. This is particularly useful when you need to sort data but retain the original indexes. It’s a way to know where each element would place if the array were sorted, without actually rearranging the array. The result is an array of indices that represent the sorted order.
Example 1: Basic Usage
Let’s start with a simple example. Consider an array arr where we need to find the indices that will sort the array:
import numpy as np
arr = np.array([3, 1, 2])
print('Original array:', arr)
argsorted_indices = arr.argsort()
print('Indices after argsort():', argsorted_indices)
This will output:
Original array: [3, 1, 2]
Indices after argsort(): [1, 2, 0]
This means that if you were to rearrange arr using these indices, element 1 (the smallest element) would be first, followed by element 2, and element 3 would be last, thus sorting the array.
Example 2: Sorting Multi-Dimensional Array along an Axis
In NumPy, you can also sort multi-dimensional arrays along a specific axis. Let’s use argsort() on a 2D array along the first axis (rows):
import numpy as np
arr = np.array([[5, 2], [3, 4]])
print('Original array:\n', arr)
row_sort_indices = arr.argsort(axis=0)
print('Row-wise sorted indices:\n', row_sort_indices)
This will output:
Original array:
[[5, 2], [3, 4]]
Row-wise sorted indices:
[[1, 0], [0, 1]]
Sorting along the row (axis=0) rearranges the elements within each column based on their values. Thus, the first column would be [3, 5] and the second column [2, 4] if sorted according to the indices provided by argsort(). Note that the actual array remains unchanged; argsort() only provides the indices.
Example 3: Applying argsort() for Complex Sorting
Sometimes, you want to sort an array based on a secondary criterion. Let’s say you have a list of employees and their corresponding salaries, and you want to sort the employees by salary:
import numpy as np
employees = np.array(['John', 'Jane', 'Doe', 'Alice'])
salaries = np.array([70000, 90000, 60000, 80000])
sorted_indices = salaries.argsort()
print('Employees sorted by salary:', employees[sorted_indices])
This will output:
Employees sorted by salary: ['Doe', 'John', 'Alice', 'Jane']
In this example, argsort() was used to sort the salaries array, and then the indices were used to rearrange the employees array. This way, you can sort data based on related arrays, preserving the relationship between them.
Example 4: Using argsort() with Structured Arrays
NumPy supports structured arrays, which let you handle complex nested arrays. Here’s how you can use argsort() to sort such an array:
import numpy as np
data = np.array([('John', 28, 70000), ('Jane', 32, 90000), ('Doe', 45, 60000), ('Alice', 38, 80000)], dtype=[('name', 'U10'), ('age', 'i4'), ('salary', 'i4')])
sorted_indices = data['salary'].argsort()
print('Data sorted by salary:\n', data[sorted_indices])
This example illustrates sorting a structured array by salary. Notice how argsort() is applied to a specific field of the structured array, in this case, ‘salary’, to obtain the sorting indices.
Conclusion
Through these examples, we’ve seen how versatile and powerful the argsort() method is for sorting and organizing data in NumPy arrays. By mastering these techniques, you can effectively manipulate and analyze your data, whether it’s simple one-dimensional lists or complex structured arrays.
