import numpy as np
%time x = np.random.randn(2 ** 26)
CPU times: user 3.86 s, sys: 72 ms, total: 3.94 s Wall time: 3.9 s
i = 1000
%%time
result = np.argsort(x)[:i]
CPU times: user 20 s, sys: 72 ms, total: 20.1 s Wall time: 20.1 s
%%time
result = np.argsort(
np.partition(x, kth=i)[:i])
CPU times: user 916 ms, sys: 60 ms, total: 976 ms Wall time: 976 ms