Returning the N largest values' indices in a multidimensional array (can find solutions for one dimension but not multi-dimension)

user_7351362 0
Viewed 982

I have a numpy array X, and I'd like to return another array Y whose entries are the indices of the n largest values of X i.e. suppose I have:

a =np.array[[1, 3, 5], [4, 5 ,6], [9, 1, 7]]

then say, if I want the first 5 "maxs"'s indices-here 9, 7 , 6 , 5, 5 are the maxs, and their indices are:

b=np.array[[2, 0], [2 2], [ 2 1],  [1 1], [0 , 2])

I've been able to find some solutions and make this work for a one dimensional array like 

c=np.array[1, 2, 3, 4, 5, 6]:

def f(a,N): 
    return np.argsort(a)[::-1][:N]

But have not been able to generate something that works in more than one dimension. Thanks!

multidimensional-array python-3.x numpy sorting
1 Answers

Approach #1

Get the argsort indices on its flattened version and select the last N indices. Then, get the corresponding row and column indices -

N = 5
idx = np.argsort(a.ravel())[-N:][::-1] #single slicing: `[:N-2:-1]`
topN_val = a.ravel()[idx]
row_col = np.c_[np.unravel_index(idx, a.shape)]

Sample run -

# Input array
In [39]: a = np.array([[1,3,5],[4,5,6],[9,1,7]])

In [40]: N = 5
    ...: idx = np.argsort(a.ravel())[-N:][::-1]
    ...: topN_val = a.ravel()[idx]
    ...: row_col = np.c_[np.unravel_index(idx, a.shape)]
    ...: 

In [41]: topN_val
Out[41]: array([9, 7, 6, 5, 5])

In [42]: row_col
Out[42]: 
array([[2, 0],
       [2, 2],
       [1, 2],
       [1, 1],
       [0, 2]])

Approach #2

For performance, we can use np.argpartition to get top N indices without keeping sorted order, like so -

idx0 = np.argpartition(a.ravel(), -N)[-N:]

To get the sorted order, we need one more round of argsort -

idx = idx0[a.ravel()[idx0].argsort()][::-1]
user_32938810
Related