Given a 3 dimensional numpy array, how to find the indexes of top n smallest values ? The index of the minimum value can be found as:
i,j,k = np.where(my_array == my_array.min())
-
2Possible duplicate of Get the indices of N highest values in an ndarraystamaimer– stamaimer2017-08-15 09:26:48 +00:00Commented Aug 15, 2017 at 9:26
-
Did the posted solution work for you?Divakar– Divakar2017-08-18 14:50:53 +00:00Commented Aug 18, 2017 at 14:50
-
Yes, it did worked. Thanks a lot matechupa_kabra– chupa_kabra2017-09-01 09:38:41 +00:00Commented Sep 1, 2017 at 9:38
Add a comment
|
1 Answer
Here's one approach for generic n-dims and generic N smallest numbers -
def smallestN_indices(a, N):
idx = a.ravel().argsort()[:N]
return np.stack(np.unravel_index(idx, a.shape)).T
Each row of the the 2D output array would hold the indexing tuple that corresponds to one of the smallest array numbers.
We can also use argpartition, but that might not maintain the order. So, we need a bit more additional work with argsort there -
def smallestN_indices_argparitition(a, N, maintain_order=False):
idx = np.argpartition(a.ravel(),N)[:N]
if maintain_order:
idx = idx[a.ravel()[idx].argsort()]
return np.stack(np.unravel_index(idx, a.shape)).T
Sample run -
In [141]: np.random.seed(1234)
...: a = np.random.randint(111,999,(2,5,4,3))
...:
In [142]: smallestN_indices(a, N=3)
Out[142]:
array([[0, 3, 2, 0],
[1, 2, 3, 0],
[1, 2, 2, 1]])
In [143]: smallestN_indices_argparitition(a, N=3)
Out[143]:
array([[1, 2, 3, 0],
[0, 3, 2, 0],
[1, 2, 2, 1]])
In [144]: smallestN_indices_argparitition(a, N=3, maintain_order=True)
Out[144]:
array([[0, 3, 2, 0],
[1, 2, 3, 0],
[1, 2, 2, 1]])
Runtime test -
In [145]: a = np.random.randint(111,999,(20,50,40,30))
In [146]: %timeit smallestN_indices(a, N=3)
...: %timeit smallestN_indices_argparitition(a, N=3)
...: %timeit smallestN_indices_argparitition(a, N=3, maintain_order=True)
...:
10 loops, best of 3: 97.6 ms per loop
100 loops, best of 3: 8.32 ms per loop
100 loops, best of 3: 8.34 ms per loop
