按多列对 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 包括一个本地 function 用于按列进行子排序， 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()

output（两者）：

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

Answer 2

您可以使用np.lexsort对多列上的数组进行排序：

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

a[idx]

Output：

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