how to get the index of the largest n values in a multi-dimensional numpy array [duplicate]

Question

I want to get the index of the largest n values in a multi-dimensional numpy array. For get the index of the largest n values in a one-dimensional numpy array, i found this. After test in interactive shell in python, it seems that bottleneck.argpartsort can't effect on multi-dimensional numpy array. For get the index of the largest value in a multi-dimensional numpy array, i found this. It can't get the largest n. The method that i can give is translate the multi-dimensional numpy array to a list of {value:index}(index present by a tuple), and then sort the list by the value, and get the index for it. Is there anything more easier or more performance?

Answer 1

I don't have access to bottleneck, so in this example I am using argsort, but you should be able to 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])