How to get indices of N minimum values per row in a NumPy matrix with Python?

Question

I have the following NumPy matrix:

m = np.array([[1, 2, 3, 4],
              [10, 5, 3, 4],
              [12, 8, 1, 2],
              [7, 0, 2, 4]])

Now, I need the indices of N (say, N=2) lowest values of each row in this matrix . So with the example above, I expect the following output:

[[0, 1],
 [2, 3],
 [3, 2],
 [1, 2]]

where the rows of the output matrix correspond to the respective rows of the original, and the elements of the rows of the output matrix are the indices of the N lowest values in the corresponding original rows (preferably in ascending order by values in the original matrix). How could I do it in NumPy?

joni · Accepted Answer · 2022-10-15 07:40:02Z

2

You could either use a simple loop-approach (not recommended) or you use np.argpartition:

In [13]: np.argpartition(m, 2)[:, :2]
Out[13]:
array([[0, 1],
       [2, 3],
       [2, 3],
       [1, 2]])

answered Oct 15, 2022 at 7:40

joni

7,1872 gold badges16 silver badges23 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Fredrik Over a year ago

Sorry, what is "2" as the second parameter of argpartition?

Rabinzel · Accepted Answer · 2022-10-15 07:39:38Z

0

You could use np.argsort on your array and then slice the array with the amount of N lowest/highest values.

np.argsort(m, axis=1)[:, :2]

array([[0, 1],
       [2, 3],
       [2, 3],
       [1, 2]], dtype=int64)

answered Oct 15, 2022 at 7:39

Rabinzel

7,9533 gold badges12 silver badges31 bronze badges

Comments

Sachin Kohli · Accepted Answer · 2022-10-15 07:39:48Z

0

Try this;

import numpy as np
m = np.array([[1, 2, 3, 4],
              [10, 5, 3, 4],
              [12, 8, 1, 2],
              [7, 0, 2, 4]])

for arr in m:
    print(arr.argsort()[:2])

answered Oct 15, 2022 at 7:39

Sachin Kohli

1,9961 gold badge4 silver badges6 bronze badges

Collectives™ on Stack Overflow

How to get indices of N minimum values per row in a NumPy matrix with Python?

3 Answers 3

1 Comment

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related