以与 sorted(, key=) 相同的方式对 numpy 数组进行排序答案

【问题标题】：Sorting numpy array in the same manner as sorted(, key=)以与 sorted(, key=) 相同的方式对 numpy 数组进行排序
【发布时间】：2013-01-18 21:58:41
【问题描述】：

例如，先按实部对复数数组进行排序，然后按虚部大小，然后再按负虚部：

def order(a):
    return a.real, abs(a.imag), sign(a.imag)

z = array(sorted(z, key=order))

所以

array([ 1.+2.j, 5.+0.j, 1.+0.j, 1.+1.j, 1.+1.j, 1.-1.j, 6.+0.j, 1.-1.j, 1.-2.j])

变成

array([ 1.+0.j, 1.-1.j, 1.-1.j, 1.+1.j, 1.+1.j, 1.-2.j, 1.+2.j, 5.+0.j, 6.+0.j])

我认为有一种方法可以使用 numpy 的 argsort 来做同样的事情，这可能更快，但我想不通：

In [2]: argsort((a.real, abs(a.imag), sign(a.imag)))
Out[2]: 
array([[0, 2, 3, 4, 5, 7, 8, 1, 6],
       [1, 2, 6, 3, 4, 5, 7, 0, 8],
       [5, 7, 8, 1, 2, 6, 0, 3, 4]])

【问题讨论】：

标签： python sorting numpy

【解决方案1】：

您可以使用np.lexsort：

import numpy as np
a = np.array([ 1.+2.j, 5.+0.j, 1.+0.j, 1.+1.j, 1.+1.j, 1.-1.j, 6.+0.j,
               1.-1.j, 1.-2.j])
sorted_idx = np.lexsort((np.sign(a.imag), np.abs(a.imag), a.real))
>>> a[sorted_idx]
array([ 1.+0.j,  1.-1.j,  1.-1.j,  1.+1.j,  1.+1.j,  1.-2.j,  1.+2.j,
        5.+0.j,  6.+0.j])

请注意，排序键的顺序是相反的，即最后一个是主体。

【讨论】：