Sorting numpy array according to the sum

Question

I want to sort an numpy array according to the sum. Something like

import numpy as np
a = np.array([1,2,3,8], [3,0,2,1])
b = np.sum(a, axis = 0)
idx = b.argsort()

Now np.take(a, idx) leads to [2, 1, 3, 8].

But I would like an array: result = np.array([2, 1, 3, 8], [0, 3, 2, 1]]

What is the most clever and fastest way to do that?

Answer 1

With the same code from your question, you can just use the optional axis argument for np.take (default the flattened array is used, that's the reason you got only the first row, see documentation):

>>> np.take(a, idx, axis=1)
array([[2, 1, 3, 8],
       [0, 3, 2, 1]])

Or you can use fancy indexing:

>>> a[:,idx]
array([[2, 1, 3, 8],
       [0, 3, 2, 1]])