Sort a numpy 2d array by 1st row, maintaining columns

In python, I have a numpy array of the form:

[4 8 2 0 5]
[3 1 6 8 1]
[2 2 6 0 3]
[9 7 6 7 8]
[5 8 1 1 4]

I want to sort it by the value of the first row from left to right in ascending order, while keeping the columns as a whole intact. The actual arrays are of unspecified dimensions, and pretty gigantic, so writing something myself with for loops get prohibitively slow. The result should be:

[0 2 4 5 8]
[8 6 3 1 1]
[0 6 2 3 2]
[7 6 9 8 7]
[1 1 5 4 8]

I can get a row vector with the column indexes ordered correctly using argsort, but don't know where to go from there on actually building the new array.

1 answer

  • answered 2018-03-20 00:05 MaxU

    Source array:

    In [215]: a
    Out[215]:
    array([[4, 8, 2, 0, 5],
           [3, 1, 6, 8, 1],
           [2, 2, 6, 0, 3],
           [9, 7, 6, 7, 8],
           [5, 8, 1, 1, 4]], dtype=int64)
    

    Using Numpy indexing:

    In [218]: a[:, a[0].argsort()]
    Out[218]:
    array([[0, 2, 4, 5, 8],
           [8, 6, 3, 1, 1],
           [0, 6, 2, 3, 2],
           [7, 6, 9, 8, 7],
           [1, 1, 5, 4, 8]], dtype=int64)
    

    Using Pandas:

    In [212]: pd.DataFrame(a).sort_values(0, axis=1).values
    Out[212]:
    array([[0, 2, 4, 5, 8],
           [8, 6, 3, 1, 1],
           [0, 6, 2, 3, 2],
           [7, 6, 9, 8, 7],
           [1, 1, 5, 4, 8]], dtype=int64)