如何使用numpy实现topk函数操作并排序
导读：本文共2319字符，通常情况下阅读需要8分钟。同时您也可以点击右侧朗读，来听本文内容。按键盘←（左） →（右） 方向键可以翻页。
摘要： topK算法经常在各种功能中使用到，在python中，numpy等计算库使用了丰富的底层优化，对于矩阵计算的效率远高于python的for-loop实现。但是numpy中并没有直接提供topK算法的直接实现。pytorch 库提供了topk函数，可以将高维数组沿某一维度（该维度共N项），选出最大（最小）的K项并排序。返回排序结果和index信息。奇怪的是，更轻量... ...
目录
(为您整理了一些要点),点击可以直达。
topK算法经常在各种功能中使用到，在python中，numpy等计算库使用了丰富的底层优化，对于矩阵计算的效率远高于python的for-loop实现。但是numpy中并没有直接提供topK算法的直接实现。
pytorch 库提供了topk函数，可以将高维数组沿某一维度（该维度共N项），选出最大（最小）的K项并排序。返回排序结果和index信息。奇怪的是，更轻量级的numpy库并没有直接提供 topK 函数。numpy只提供了argpartition 和 partition，可以将最大（最小）的K项排到前K位。以argpartition为例，最小的3项排到了前3位：
注意，argpartition实现的是 partial sorting，如上例，前3项和其余项被分开，但是两部分各自都是不排序的！而我们可能更想要topK的几项排好序（其余项则不作要求）。因此，下面提供一种基于argpartition的topK方法。
最简单的方法自然是全排序，然后取前K项。缺点在于，要把topK之外的数据也进行排序，当K << N时较为浪费时间，复杂度为O ( n log n ) O(n log n)O(nlogn):
对于 np.argpartition 函数，复杂度可能下降到 O ( n log K ) O(n log K)O(nlogK)，很多情况下，K << N，此时naive方法有优化的空间。
以下方法首先选出 topK 项，然后仅对前topK项进行排序（matrix仅限2d-array）。
对shape(5000, 100000)的矩阵进行topK排序，测试时间为：
补充：python堆排序实现TOPK问题
如何使用numpy实现topk函数操作并排序的详细内容，希望对您有所帮助，信息来源于网络。