这篇教程使用numpy实现topk函数操作(并排序)写得很实用，希望能帮到您。
np.argpartition 难以解决topKtopK是常用的一个功能，在python中，numpy等计算库使用了丰富的底层优化，对于矩阵计算的效率远高于python的for-loop实现。因此，我们希望尽量用一些numpy函数的组合实现topK。 pytorch 库提供了topk函数，可以将高维数组沿某一维度（该维度共N项），选出最大（最小）的K项并排序。返回排序结果和index信息。奇怪的是，更轻量级的numpy库并没有直接提供 topK 函数。numpy只提供了argpartition 和 partition，可以将最大（最小）的K项排到前K位。以argpartition为例，最小的3项排到了前3位： >>> x = np.array([3, 5, 6, 4, 2, 7, 1])>>> x[np.argpartition(x, 3)]array([2, 1, 3, 4, 5, 7, 6]) 注意，argpartition实现的是 partial sorting，如上例，前3项和其余项被分开，但是两部分各自都是不排序的！而我们可能更想要topK的几项排好序（其余项则不作要求）。因此，下面提供一种基于argpartition的topK方法。 一个naive方法最简单的方法自然是全排序，然后取前K项。缺点在于，要把topK之外的数据也进行排序，当K << N时较为浪费时间，复杂度为O ( n log n ) O(n /log n)O(nlogn): def naive_arg_topK(matrix, K, axis=0): """ perform topK based on np.argsort :param matrix: to be sorted :param K: select and sort the top K items :param axis: dimension to be sorted. :return: """ full_sort = np.argsort(matrix, axis=axis) return full_sort.take(np.arange(K), axis=axis)# Example>>> dists = np.random.permutation(np.arange(30)).reshape(6, 5)array([[17, 28, 1, 24, 23, 8], [ 9, 21, 3, 22, 4, 5], [19, 12, 26, 11, 13, 27], [10, 15, 18, 14, 7, 16], [ 0, 25, 29, 2, 6, 20]])>>> naive_arg_topK(dists, 2, axis=0)array([[4, 2, 0, 4, 1, 1], [1, 3, 1, 2, 4, 0]])>>> naive_arg_topK(dists, 2, axis=1)array([[2, 5], [2, 4], [3, 1], [4, 0], [0, 3]]) 基于partition的方法对于 np.argpartition 函数，复杂度可能下降到 O ( n log K ) O(n /log K)O(nlogK)，很多情况下，K << N，此时naive方法有优化的空间。 以下方法首先选出 topK 项，然后仅对前topK项进行排序（matrix仅限2d-array）。 def partition_arg_topK(matrix, K, axis=0): """ perform topK based on np.argpartition :param matrix: to be sorted :param K: select and sort the top K items :param axis: 0 or 1. dimension to be sorted. :return: """ a_part = np.argpartition(matrix, K, axis=axis) if axis == 0: row_index = np.arange(matrix.shape[1 - axis]) a_sec_argsort_K = np.argsort(matrix[a_part[0:K, :], row_index], axis=axis) return a_part[0:K, :][a_sec_argsort_K, row_index] else: column_index = np.arange(matrix.shape[1 - axis])[:, None] a_sec_argsort_K = np.argsort(matrix[column_index, a_part[:, 0:K]], axis=axis) return a_part[:, 0:K][column_index, a_sec_argsort_K]# Example>>> dists = np.random.permutation(np.arange(30)).reshape(6, 5)array([[17, 28, 1, 24, 23, 8], [ 9, 21, 3, 22, 4, 5], [19, 12, 26, 11, 13, 27], [10, 15, 18, 14, 7, 16], [ 0, 25, 29, 2, 6, 20]])>>> partition_arg_topK(dists, 2, axis=0)array([[4, 2, 0, 4, 1, 1], [1, 3, 1, 2, 4, 0]])>>> partition_arg_topK(dists, 2, axis=1)array([[2, 5], [2, 4], [3, 1], [4, 0], [0, 3]]) 大数据量测试对shape(5000, 100000)的矩阵进行topK排序，测试时间为： | K | partition(s) | naive(s) | | 10 | 8.884 | 22.604 | | 100 | 9.012 | 22.458 | | 1000 | 8.904 | 22.506 | | 5000 | 11.305 | 22.844 | 
补充：python堆排序实现TOPK问题 
# 构建小顶堆跳转def sift(li, low, higt): tmp = li[low] i = low j = 2 * i + 1 while j <= higt: # 情况2：i已经是最后一层 if j + 1 <= higt and li[j + 1] < li[j]: # 右孩子存在并且小于左孩子 j += 1 if tmp > li[j]: li[i] = li[j] i = j j = 2 * i + 1 else: break # 情况1：j位置比tmp小 li[i] = tmpdef top_k(li, k): heap = li[0:k] # 建堆 for i in range(k // 2 - 1, -1, -1): sift(heap, i, k - 1) for i in range(k, len(li)): if li[i] > heap[0]: heap[0] = li[i] sift(heap, 0, k - 1) # 挨个输出 for i in range(k - 1, -1, -1): heap[0], heap[i] = heap[i], heap[0] sift(heap, 0, i - 1) return heapli = [0, 8, 6, 2, 4, 9, 1, 4, 6]print(top_k(li, 3)) 以上为个人经验，希望能给大家一个参考，也希望大家多多支持51zixue.net。
像线程一样管理进程的Python multiprocessing库
Python统计可散列的对象之容器Counter详解 |
