python - Find the index of the n smallest values in a 3 dimensional numpy array -
given 3 dimensional numpy array, how find indexes of top n smallest values ? index of minimum value can found as:
i,j,k = np.where(my_array == my_array.min())
here's 1 approach generic n-dims , 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 2d
output array hold indexing tuple corresponds 1 of smallest array numbers.
we can use argpartition
, might not maintain order. so, need bit more additional work 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) ...: = 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]: = 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
Comments
Post a Comment