如何在Python中对形状为(m,n)的NumPy数组按列分层排序？

要实现对(m,n)形状的NumPy数组按列进行分层排序（即先按第一列排序，再按第二列排序，以此类推），可以直接使用NumPy内置的np.lexsort()函数，该函数能高效完成多键排序，完全满足速度要求。

实现步骤 #

• 准备输入数组：将目标数据转换为NumPy数组格式。
• 生成排序索引：使用np.lexsort()传入反转后的列作为排序键（lexsort会按传入键的逆序优先排序，最后一个传入的键是主排序键）。
• 应用索引得到结果：用生成的索引对原数组切片，得到分层排序后的数组。

代码示例 #

import numpy as np

# 输入数组
arr = np.array([
    [2, 6],
    [1, 7],
    [1, 1],
    [2, 1],
    [1, 4],
    [1, 1]
])

# 生成分层排序索引：优先按第一列排序，再按第二列排序
sorted_indices = np.lexsort(arr.T[::-1])

# 应用索引得到排序后的数组
sorted_arr = arr[sorted_indices]

print(sorted_arr)

输出结果 #

[[1 1]
 [1 1]
 [1 4]
 [1 7]
 [2 1]
 [2 6]]

说明 #

• np.lexsort()是稳定排序，会保留相同键值元素的原始相对顺序。
• 对于任意n列的数组，只需保持arr.T[::-1]的传入方式，即可自动按从第一列到第n列的优先级完成分层排序。

内容的提问来源于stack exchange，提问作者peets