How to get the index of the largest n values ​​in a numpy multidimensional array

Question

I want to get the index of the largest n values ​​in a multidimensional numpy array. To get the index of the largest n values ​​in a one-dimensional numpy array, I found this . After the test in the interactive shell in python, it seems that bottleneck.argpartsortit cannot affect the numpy multidimensional array. To get the index of the highest value in a numpy multidimensional array, I found this . He cannot get the greatest n. The method I can give translates the numpy multidimensional array into a list {value:index}(the index present by the tuple), and then sorts the list by value and gets the index for it. Is there something simpler or more efficient?

Answer 1

I do not have access to bottleneck, so in this example I use argsort, but you can use it in the same way:

#!/usr/bin/env python
import numpy as np
N = 4
a = np.random.random(20).reshape(4, 5)
print(a)

# Convert it into a 1D array
a_1d = a.flatten()

# Find the indices in the 1D array
idx_1d = a_1d.argsort()[-N:]

# convert the idx_1d back into indices arrays for each dimension
x_idx, y_idx = np.unravel_index(idx_1d, a.shape)

# Check that we got the largest values.
for x, y, in zip(x_idx, y_idx):
    print(a[x][y])