numpy中的argpartition用法

毕业论文写完了,闲来无事搞了一个汉字书法识别比赛,提交的结果需要是Top-5,在google排序方法的时候,发现一个比argsort效率更高的方法:argpartition

$ import numpy as np
$ a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
$ print(np.argpartition(a, 4)) #将数组a中所有元素(包括重复元素)从小到大排列,比第5大的元素小的放在前面,大的放在后面,输出新数组索引
>> [6 9 4 3 7 2 1 5 8 0]
$ a[np.argpartition(a, 4)]     #输出新数组索引对应的数组
>> array([0, 0, 3, 3, 4, 4, 4, 9, 6, 9])

注意,排序规则是从0开始排序,即,第0大的元素为“0”,第1大的元素还为“0”,第2大的元素为“3”,以此类推。(我试了好几次才发现的。。。)

如果想输出Top5:

$ import numpy as np
$ a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
$ a[np.argpartition(a, -5)[-5:]]
>> array([4, 4, 9, 6, 9])

猜你喜欢

转载自blog.csdn.net/MacwinWin/article/details/80002584
今日推荐