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

Admin 1
Viewed 108

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!

from his mujjjga's answer i was able to fine one array that does not work

np.array([100,100,100,100,111,111,90],
[102,102,102,102,102,102,102],
[150,150,150,150,95,95,95]])
numpy sorting arrays python
1 Answers

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

Answer updated to fix the bug:

The last row is the primary key traversing from the last row towards first to handle ties

def lexsort2D_copy(data):
  idxs = np.arange(data.shape[1])
  for d in data:
    idxs = idxs[np.argsort(d[idxs], kind="stable")]
  return idxs

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]])
assert np.all(np.lexsort(test) == lexsort2D_copy(test))

test=np.array([[100,100,100,100,111,400,120],
               [229,1133,152,210,120,320,320]])
assert np.all(np.lexsort(test) == lexsort2D_copy(test))
 
test = np.array([[100,100,100,100,111,111,90],
                [ 102,102,102,102,102,102,102],
                [ 150,150,150,150,95,95,95]])

assert np.all(np.lexsort(test) == lexsort2D_copy(test))
user_4239260
Related