how to generate same result using 2D array of np.argsort to lexsort?

Question

I've replicated using 1d array argsort that can matches with lexsort.

#a = 1d np.array
#b = 1d np.array

def lexsort_copy(a,b):
    idxs= np.argsort(a,kind='stable')
    return idxs[np.argsort(b[idxs],kind='stable')]

lexsort_copy(a,b) == np.lexsort((a,b))

which gives me the same output, but I am struggling how to replicate this using 2d array.

test 2d array:

test=np.array([[100,100,100,100,111,400,120],[229,1133,152,210,120,320,320]])
np.lexsort(test)

output:

array([4, 2, 3, 0, 6, 5, 1], dtype=int64)

how can we replicate this above output without using lexsort for 2d array?

Any solution here would be appreciated! Thank you!

Answer 1

You can extend your lexsort_copy to work with 2d arrays as below:

def lexsort2D_copy(data):
  idxs = np.arange(data.shape[1])
  for i in range(data.shape[0]-1):
    idxs = np.argsort(data[i][idxs],kind='stable')
  return idxs[np.argsort(data[-1][idxs],kind='stable')]

Test:

test=np.array([[100,100,100,100,111,400,120],
               [229,1133,152,210,120,320,320],
               [29,133,12,10,10,20,3120]])
np.lexsort(test)  == lexsort2D_copy(test)

test=np.array([[100,100,100,100,111,400,120],
               [229,1133,152,210,120,320,320]])
np.lexsort(test)  == lexsort2D_copy(test)

Output:

array([ True,  True,  True,  True,  True,  True,  True])
array([ True,  True,  True,  True,  True,  True,  True])