按多列对 numpy 二维数组进行排序

Question

我有一个 2D numpy 数组，如下所示

array([[5, 0],
       [3, 1],
       [7, 0],
       [2, 1]])

我想按每列（从右到左）进行（子）排序以获得以下结果：

array([[5, 0],
       [7, 0],
       [2, 1],
       [3, 1]])

我怎样才能在 numpy 中做到这一点？

Answer 1

Numpy 包含一个用于按列进行子排序的本机函数lexsort：

idx = np.lexsort((arr[:,0], arr[:,1]))
arr_sorted = arr[idx]

或者，如果您更熟悉，也可以使用 pandas 语法；这会产生一些内存/时间开销，但对于 < 1m 行来说应该很小：

arr = [
    [5,  0],
    [3,  1],
    [7,  0],
    [2,  1]
]
df = pd.DataFrame(data=arr).sort_values([1,0])
arr_sorted = df.to_numpy()

输出（两者）：

array([[5, 0],
       [7, 0],
       [2, 1],
       [3, 1]])