插件窝 干货文章 NumPy数组排序:如何高效地按指定列排序?

NumPy数组排序:如何高效地按指定列排序?

数组 排序 索引 sorted 937    来源:    2025-03-16

在NumPy中,你可以使用 numpy.argsort() 函数来高效地按指定列对数组进行排序。argsort() 返回的是数组排序后的索引,你可以利用这些索引来重新排列数组。

假设你有一个二维数组 arr,你想按第 k 列进行排序,可以按照以下步骤操作:

  1. 使用 argsort() 获取排序后的索引。
  2. 使用这些索引来重新排列数组。

以下是一个示例代码:

import numpy as np

# 创建一个示例二维数组
arr = np.array([[3, 2, 5],
                [1, 4, 6],
                [7, 0, 8]])

# 指定要排序的列(例如第1列,索引为0)
k = 0

# 获取排序后的索引
sorted_indices = np.argsort(arr[:, k])

# 使用索引重新排列数组
sorted_arr = arr[sorted_indices]

print("排序后的数组:")
print(sorted_arr)

解释:

  • arr[:, k]:获取第 k 列的所有元素。
  • np.argsort(arr[:, k]):返回第 k 列排序后的索引。
  • arr[sorted_indices]:使用这些索引重新排列整个数组。

输出:

排序后的数组:
[[1 4 6]
 [3 2 5]
 [7 0 8]]

按多列排序

如果你想按多列排序,可以使用 numpy.lexsort() 函数。lexsort() 允许你按多个键进行排序。

# 按第1列和第2列排序
sorted_indices = np.lexsort((arr[:, 1], arr[:, 0]))

# 使用索引重新排列数组
sorted_arr = arr[sorted_indices]

print("按第1列和第2列排序后的数组:")
print(sorted_arr)

输出:

按第1列和第2列排序后的数组:
[[1 4 6]
 [3 2 5]
 [7 0 8]]

通过这些方法,你可以高效地对NumPy数组按指定列进行排序。