Is there a way to use np.min to get more than 1 minimum number from matrix?
x = np.array([[1,2,3],[4,5,6],[7,8,0]])
Expected result for 3 minimum numbers
>>[0,1,2]
Add a comment
|
2 Answers
One way would be to use argsort on flattend view -
x.ravel()[x.ravel().argsort()[:3]]
More performant one with np.argpartition -
x.ravel()[x.ravel().argpartition(range(3))[:3]]
Or with sort to sort it afterwards -
np.sort(x.ravel()[x.ravel().argpartition(3)[:3]])
If you don't care about the numbers being sorted, skip sort -
x.ravel()[x.ravel().argpartition(3)[:3]]
Sample run -
In [44]: x
Out[44]:
array([[1, 2, 3],
[4, 5, 6],
[7, 8, 0]])
In [45]: x.ravel()[x.ravel().argsort()[:3]]
Out[45]: array([0, 1, 2])
In [48]: x.ravel()[x.ravel().argpartition(range(3))[:3]]
Out[48]: array([0, 1, 2])
In [52]: np.sort(x.ravel()[x.ravel().argpartition(3)[:3]])
Out[52]: array([0, 1, 2])
In [47]: x.ravel()[x.ravel().argpartition(3)[:3]]
Out[47]: array([0, 1, 2])
1 Comment
piRSquared
argpartition is the way to go. Its O(log n) where argsort is O(n log n)Flatten and sort the array and get first 3 elements.
sorted(x.flatten())[:3]
Out[275]: [0, 1, 2]
Or a faster approach:
sorted(sum(x.tolist(),[]))[:3]
%timeit sorted(sum(x.tolist(),[]))[:3]
The slowest run took 7.38 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 2.14 µs per loop
%timeit sorted(x.flatten())[:3]
The slowest run took 11.54 times longer than the fastest. This could mean that an intermediate result is being cached.
100000 loops, best of 3: 4.97 µs per loop
