Sort matrix columns based on the values in the first row

If matrix is your input, you can do:

result = list(zip(*sorted(zip(*matrix))))

So working from inside out, this expression does:

• zip: to iterate the transposed of the matrix (rows become columns and vice versa)
• sorted: sorts the transposed matrix. No need to provide a custom key, the sorting will be by the first element (row, which is a column in the original matrix). If there is a tie, by second element (row), ...etc.
• zip: to iterate the transposed of the transposed matrix, i.e. transposing it back to its original shape
• list to turn the iterable to a list (a matrix)

The implementation here can be very simple depending on how the data, ie. the matrix, is represented. If it is given as a list of column-lists, it just needs a sort. For the given example:

>>> m = [[2, 3, 7], [-1, -2, 5.2], [0, 1, 4], [2, 4, 5]]
>>> y = sorted(m, key=lambda x: x[0])
>>> y
[[-1, -2, 5.2], [0, 1, 4], [2, 3, 7], [2, 4, 5]]

Other representations might need a more complex approach. For example, if the matrix is given as a list of rows:

>>> m = [[2, -1, 0, 2], [3, -2, 1, 4], [7, 5.2, 4, 5]]
>>> order = sorted(range(len(m[0])), key=lambda x: m[0][x])
>>> order
[1, 2, 0, 3]
>>> y = [[row[x] for x in order] for row in m]
>>> y
[[-1, 0, 2, 2], [-2, 1, 3, 4], [5.2, 4, 7, 5]]

The idea here is that first, we will get the order the elements are going to be in based on the first row. We can do that by sorting range(4), so [0, 1, 2, 3] with the sorting key (the value used for sorting) being the i-th value of the first row.

The result is that we get [1, 2, 0, 3] which says: Column index 1 is first, then index 2, 0 and finally 3.

Now we want to create a new matrix where every row follows that order which we can do with a list comprehension over the original matrix, where for each row, we create a new list that has the elements of the row according to the order we determined before.

Note that this approach creates new lists for the whole matrix - if you're dealing with large matrices, you probably want to use the appropriate primitives from numpy and swap the elements around in place.