按多个轴排序2D numpy数组

Question

我有一个2D numpy形状的阵列(N,2),它持有N个点(x和y坐标).例如:

array([[3, 2],
       [6, 2],
       [3, 6],
       [3, 4],
       [5, 3]])

我想对它进行排序,使得我的点按x坐标排序,然后在x坐标相同的情况下按y排序.所以上面的数组应该如下所示:

array([[3, 2],
       [3, 4],
       [3, 6],
       [5, 3],
       [6, 2]])

如果这是一个普通的Python列表,我只需要定义一个比较器来做我想要的,但据我所知,numpy的sort函数不接受用户定义的比较器.有任何想法吗？

---

编辑:感谢您的想法!我设置了一个包含1000000个随机整数点的快速测试用例,并对我可以运行的那些进行了基准测试(抱歉,目前无法升级numpy).

Mine:   4.078 secs 
mtrw:   7.046 secs
unutbu: 0.453 secs

Answer 1

使用lexsort:

import numpy as np    
a = np.array([(3, 2), (6, 2), (3, 6), (3, 4), (5, 3)])

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

a[ind]
# array([[3, 2],
#       [3, 4],
#       [3, 6],
#       [5, 3],
#       [6, 2]])

a.ravel()如果a是,则返回视图C_CONTIGUOUS.如果这是真的, @ ars的方法,通过使用ravel而不是稍微修改flatten,产生一个很好的方法来a 就地排序:

a = np.array([(3, 2), (6, 2), (3, 6), (3, 4), (5, 3)])
dt = [('col1', a.dtype),('col2', a.dtype)]
assert a.flags['C_CONTIGUOUS']
b = a.ravel().view(dt)
b.sort(order=['col1','col2'])

由于b是一个视图a,排序b排序a,以及:

print(a)
# [[3 2]
#  [3 4]
#  [3 6]
#  [5 3]
#  [6 2]]

Answer 2

标题写着"排序2D数组".虽然提问者使用了一个(N,2)形状的数组,但是可以将unutbu的解决方案推广到任何(N,M)数组,因为这是人们可能真正想要的.

可以transpose使用数组并使用带有负数的切片表示法以相反的顺序step传递所有列lexsort:

>>> import numpy as np
>>> a = np.random.randint(1, 6, (10, 3))
>>> a
array([[4, 2, 3],
       [4, 2, 5],
       [3, 5, 5],
       [1, 5, 5],
       [3, 2, 1],
       [5, 2, 2],
       [3, 2, 3],
       [4, 3, 4],
       [3, 4, 1],
       [5, 3, 4]])

>>> a[np.lexsort(np.transpose(a)[::-1])]
array([[1, 5, 5],
       [3, 2, 1],
       [3, 2, 3],
       [3, 4, 1],
       [3, 5, 5],
       [4, 2, 3],
       [4, 2, 5],
       [4, 3, 4],
       [5, 2, 2],
       [5, 3, 4]])