[파이썬][numpy] numpy `argsort` 함수

In this blog post, we will explore the argsort function in the popular numerical computing library, NumPy, for Python. argsort is a powerful tool that allows us to obtain the indices that would sort an array or a list in ascending order.

Understanding argsort

The argsort function in NumPy returns the indices that would sort the input array or list. It can be used to efficiently sort large arrays while keeping track of the original indices.

The syntax for argsort is as follows:

numpy.argsort(a, axis=-1, kind='quicksort', order=None)

Here,

Examples

Let’s dive into a few examples to illustrate the usage of argsort.

Example 1: Sorting a NumPy array

import numpy as np

arr = np.array([5, 2, 8, 1, 9])

sorted_indices = np.argsort(arr)
sorted_array = arr[sorted_indices]

print(sorted_indices)  # Output: [3 1 0 2 4]
print(sorted_array)  # Output: [1 2 5 8 9]

In this example, we have an array arr containing five numbers. We use argsort to obtain the indices that would sort the array in ascending order. By applying these indices to the original array, we get a new array sorted_array with the elements sorted in ascending order.

Example 2: Sorting a multi-dimensional NumPy array along an axis

import numpy as np

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

sorted_indices = np.argsort(arr, axis=1)
sorted_array = np.take_along_axis(arr, sorted_indices, axis=1)

print(sorted_indices)
# Output: [[1 0 2]
#          [0 2 1]]
print(sorted_array)
# Output: [[2 5 8]
#          [1 4 9]]

In this example, we have a 2D array arr and we want to sort the elements along the second axis (axis=1). We use argsort with the axis parameter set to 1 to obtain the sorted indices. Finally, we use np.take_along_axis to apply the indices and get the sorted array.

Conclusion

The argsort function in NumPy provides a convenient way to obtain the indices that would sort an array or a list in ascending order. It is a versatile tool that can be used in various scenarios, from simple sorting to more complex tasks involving multi-dimensional arrays. By understanding and utilizing argsort, you can efficiently sort and manipulate data in your Python programs.