Сортировка массива numpy таким же образом, как отсортировано (, key =)

Например, для сортировки массива комплексных чисел сначала по вещественной части, затем по величине мнимой части, а затем сначала с отрицательными мнимыми частями:

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])

Я думаю, что есть способ сделать то же самое с помощью argsort от numpy, что, вероятно, быстрее, но я не могу понять:

 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]]) 

Вы можете использовать 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]) 

Обратите внимание, что ключи сортировки находятся в обратном порядке, т. Е. Последний является основным.