对于像您对部分排序索引感兴趣的情况,有NumPy's argpartition.
你有麻烦事np.argsort
in : w[np.argsort(z)[::-1]][:7]
,这本质上是w[idx]
, where idx = np.argsort(z)[::-1][:7]
.
So, idx
可以计算为np.argpartition
,就像这样-
idx = np.argpartition(-z,np.arange(7))[:7]
That -z
需要,因为默认情况下np.argpartition
尝试按升序获取排序索引。所以,为了reverse它,我们已经否定了要素。
因此,对原始代码的拟议更改将是:
func = w[np.argpartition(-z,np.arange(7))[:7]]
运行时测试 -
In [162]: z = np.random.randint(0,10000000,(1100000)) # Random int array
In [163]: idx1 = np.argsort(z)[::-1][:7]
...: idx2 = np.argpartition(-z,np.arange(7))[:7]
...:
In [164]: np.allclose(idx1,idx2) # Verify results
Out[164]: True
In [165]: %timeit np.argsort(z)[::-1][:7]
1 loops, best of 3: 264 ms per loop
In [166]: %timeit np.argpartition(-z,np.arange(7))[:7]
10 loops, best of 3: 36.5 ms per loop