Skip to content

Instantly share code, notes, and snippets.

@bobofzhang

bobofzhang/gist:4166894

Forked from tomerfiliba/gist:3698403
Created November 29, 2012 04:57
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save bobofzhang/4166894 to your computer and use it in GitHub Desktop.
Save bobofzhang/4166894 to your computer and use it in GitHub Desktop.
Download ZIP
get the indexes of the top n elements in a numpy 2d-array
Raw
gistfile1.py
import numpy as np
import bottleneck as bn
def top_n_indexes(arr, n):
idx = bn.argpartsort(arr, arr.size-n, axis=None)[-n:]
width = arr.shape[1]
return [divmod(i, width) for i in idx]
np.random.seed(47)
arr = np.random.rand(50, 50)
idx = top_n_indexes(arr, 8)
idx.sort(key = lambda tup: tup[0])
print idx
# [(5, 45), (11, 8), (16, 14), (24, 40), (25, 35), (31, 7), (35, 6), (49, 43)]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment