Sort np array based on summed selected values of each row

July 11, 2022

I have a 2D numpy array, filled with floats.
I want to take a selected chunk of each row (say item 2nd to 3rd), sum these values and sort all the rows based on that sum in a descending order.

For example:

array([[0.80372444, 0.35468653, 0.9081662 , 0.69995566],
       [0.53712474, 0.90619077, 0.69068265, 0.73794143],
       [0.14056974, 0.34685164, 0.87505744, 0.56927803]])

Here’s what I tried:

a = np.array(sorted(a, key = sum))

But that just sums all values from each row, rather that, say, only 2nd to 6th element.

>Solution :

You can start by using take to get elements at indices [1,2] from each row (axis = 1). Then sum across those element for each row (again axis = 1), and use argsort to get the order of the sums. This gives a set of row indices, which you can use to slice the array in the desired order.

import numpy as np

a = np.array([[0.80372444, 0.35468653, 0.9081662 , 0.69995566],
              [0.53712474, 0.90619077, 0.69068265, 0.73794143],
              [0.14056974, 0.34685164, 0.87505744, 0.56927803]])

a[a.take([1, 2], axis=1).sum(axis=1).argsort()]
# returns:
array([[0.14056974, 0.34685164, 0.87505744, 0.56927803],
       [0.80372444, 0.35468653, 0.9081662 , 0.69995566],
       [0.53712474, 0.90619077, 0.69068265, 0.73794143]])