Sorting 2D array by the first n rows

Question

How can I sort an array in NumPy by the two first rows?

For example,

A=array([[9, 2, 2],
         [4, 5, 6],
         [7, 0, 5]])

And I'd like to sort columns by the first two rows, such that I get back:

A=array([[2, 2, 9],
         [5, 6, 4],
         [0, 5, 7]])

Thank you!

Answer 1

One approach is to transform the 2D array over which we want to take the argsort into an easier to handle 1D array. For that one idea could be to multiply the rows to take into accounts for the sorting purpose by successively decreasing values in the power of 10 sequence, sum them and then use argsort ( note : this method will be numerically unstable for high values of k . Meant for values up to ~ 20 ):

def sort_on_first_k_rows(x, k):
    # normalize each row so that its max value is 1
    a = (x[:k,:]/x[:k,:,None].max(1)).astype('float64')
    # multiply each row by the seq 10^n, for n=k-1,k-2...0
    # Ensures that the contribution of each row in the sorting is
    # captured in the final sum
    a_pow = (a*10**np.arange(a.shape[0]-1,-1,-1)[:,None])
    # Sort with the argsort on the resulting sum
    return x[:,a_pow.sum(0).argsort()]

Checking with the shared example:

sort_on_first_k_rows(A, 2)
array([[2, 2, 9],
       [5, 6, 4],
       [0, 5, 7]])

---

Or with another example:

A=np.array([[9, 2, 2, 1, 5, 2, 9],
            [4, 7, 6, 0, 9, 3, 3],
            [7, 0, 5, 0, 2, 1, 2]])

sort_on_first_k_rows(A, 2)
array([[1, 2, 2, 2, 5, 9, 9],
       [0, 3, 6, 7, 9, 3, 4],
       [0, 1, 5, 0, 2, 2, 7]])

Answer 2

The pandas library is very flexible for sorting DataFrames - but only based on columns. So I suggest to transpose and convert your array to a DataFrame like this (note that you need to specify column names for later defining the sorting criteria):

df = pd.DataFrame(A.transpose(), columns=['col'+str(i) for i in range(len(A))])

Then sort it and convert it back like this:

A_new = df.sort_values(['col0', 'col1'], ascending=[True, True]).to_numpy().transpose()