Numpy 按降序对二维数组进行排序并从每行中取前 N

Question

here have an original 2D array

in_arr = np.array([[20,0,10,40,30], [50,40,60,90,80]])

# original array
# [[20,  0, 10, 40, 30],
#  [50, 40, 60, 90, 80]]

I need to sort the array by descending and by row, therefore, I use numpy.argsort(axis=1), and output sorted indices I got

out_arr1 = np.argsort(in_arr, axis = 1)[:,::-1]
>>> array([[3, 4, 0, 2, 1],
          [3, 4, 2, 0, 1]])

Then, I need to extract the first 3 largerst number from each array row, the sample desired output as follows:

# first 3 largest number from each row
# [[40,30,20],
#  [90,80,60]]

I have struggling few hours to try to come out correct solution, but still have no idea how should I do, here would like to ask for help. Your valuable time and advice will be much appreciated. Thank you!

Answer 1

Using numpy.argsort() returns an array of indices for the sorted array. As such, what your out_arr1 lets you know is where on each row to find the highest values.

If you are to continue this way, what you would need to do is for each row in in_arr (hereby written as in_arr[i] ) take values found at the first 3 indices in out_arr1[i] .

What that means is that out_arr1[i, 0] tells you where the highest value in in_arr on row i is located. In our case, out_arr1[0, 0] = 3 , which means the highest value in row 0 is 40 (on index 3)

Doing this, the 3 largest numbers on each row are represented by out_arr1[0, 0] , out_arr1[0, 1] , out_arr1[0, 2] and out_arr1[1, 0] , out_arr1[1, 1] , out_arr1[1, 2] .

to get the desired output, we would need something along the lines of:

final_arr = numpy.array([in_arr[0, out_arr1[0, 0], in_arr[0, out_arr1[0, 1], in_arr[0, out_arr1[0, 2], in_arr[1, out_arr1[1, 0], in_arr[1, out_arr1[1, 1], in_arr[1, out_arr1[1, 2]])

This however, is less than elegant, and there is another, easier solution to your problem.

Using numpy.sort() instead of numpy.argsort() we can return the exact values of in_arr sorted along an axis. By doing that, we no longer need to use an output index to find our 3 highest values, as they are the first 3 in our new output.

Considering out_arr2 as the output from numpy.sort() , the final array would look like:

final_arr = numpy.array([[out_arr[0, 0], out_arr[0, 1], out_arr[0, 2]], [out_arr[1, 0], out_arr[1, 1], out_arr[1, 2]]])

Answer 2

Based on this this answer you can do something like this

np.array(list(map(lambda x, y: y[x], np.argsort(in_arr), in_arr)))[:,::-1][:,:3]

which gives

array([[40, 30, 20],
       [90, 80, 60]])

Answer 3

You can first sort all rows in the input array with a list comprehension using sorted . Then you extract the last 3 numbers of the rows.

in_arr = np.array([[20,0,10,40,30], [50,40,60,90,80]])

output = []
for i in [sorted(row) for row in in_arr]:
    output.append(i[-3:][::-1])
    
print(output)