Python
numpy

argsort and argpartition method sample in numpy

np.set_printoptions(threshold=np.inf)
Z = np.arange(10)
np.random.shuffle(Z)
print(Z)
n = 3

# Slow
a = np.argsort(Z)
print(a)
print (Z[a[-n:]])
# Fast
print(-Z)
a = np.argpartition(-Z,n)
print(a)
print (Z[a[:n]])
  • e.g.
[0 5 6 9 3 1 2 8 4 7]
[0 5 6 4 8 1 2 9 7 3]
[7 8 9]
[ 0 -5 -6 -9 -3 -1 -2 -8 -4 -7]
[3 7 9 2 1 8 4 6 5 0]
[9 8 7]
  • argpartition method results is not sorted.

Refs

https://gist.github.com/naoyashiga/8f8a215932e881a3f9ec85e45d499e99
* 省略表記無効