首页 > python教程

使用numpy实现topk函数操作(并排序)

时间:2021-07-08 python教程 查看: 889

np.argpartition 难以解决topK

topK是常用的一个功能,在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] = tmp


def 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 heap


li = [0, 8, 6, 2, 4, 9, 1, 4, 6]
print(top_k(li, 3))

以上为个人经验,希望能给大家一个参考,也希望大家多多支持python博客。

展开全文
上一篇:Keras 多次加载model出错的解决方案
下一篇:Python统计可散列的对象之容器Counter详解
输入字:
相关知识
Python 实现图片色彩转换案例

我们在看动漫、影视作品中,当人物在回忆过程中,体现出来的画面一般都是黑白或者褐色的。本文将提供将图片色彩转为黑白或者褐色风格的案例详解,感兴趣的小伙伴可以了解一下。

python初学定义函数

这篇文章主要为大家介绍了python的定义函数,具有一定的参考价值,感兴趣的小伙伴们可以参考一下,希望能够给你带来帮助,希望能够给你带来帮助

图文详解Python如何导入自己编写的py文件

有时候自己写了一个py文件,想要把它导入到另一个py文件里面,所以下面这篇文章主要给大家介绍了关于Python如何导入自己编写的py文件的相关资料,需要的朋友可以参考下

python二分法查找实例代码

二分算法是一种效率比较高的查找算法,其输入的是一个有序的元素列表,如果查找元素包含在列表中,二分查找返回其位置,否则返回NONE,下面这篇文章主要给大家介绍了关于python二分法查找的相关资料,需要的朋友可以参考下