Sorting a numpy array in the same way as sorting (, key =)

Question

For example, to sort an array of complex numbers, first the real part, then the magnitude of the imaginary part, and then first with the negative imaginary parts:

def order(a): return a.real, abs(a.imag), sign(a.imag) z = array(sorted(z, key=order)) 

So

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

becomes

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

I think there is a way to do the same using numpy argsort , which is probably faster, but I can't figure it out from:

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

Answer 1

You can use 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]) 

Note that the sort keys are in reverse order, i.e. the latter is the main one.