使用numpy实现topk函数操作(并排序)
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))
(资源库 www.zyku.net)
原文链接:https://blog.csdn.net/SoftPoeter/article/details/86629329
您可能感兴趣的文章
- 01-09Symbolab函数-Symbolab函数应用软件功能介绍
- 06-16PHP函数file_get_contents被屏蔽解决方法
- 05-31PHP常用的转义字符函数介绍
- 05-31帝国CMS二次开发经常会用的ehtmlspecialchars函数介绍
- 03-28Python 执行函数的九种方法
- 03-15微信小程序 onLoad 函数
- 03-13帝国CMS提示信息函数printerror()
- 03-13帝国CMS常用函数介绍(二次开发参考)
- 03-13帝国CMS常用函数
- 03-13帝国CMS获取信息内容页地址函数sys_ReturnBqTitleLink
- 08-30织梦DedeCMS获取文章链接的函数GetOneArchive使用方法
- 11-30$_SERVER函数中QUERY_STRING和REQUEST_URI区别详解
- 07-27PHP自定义函数判断是否为Get、Post及Ajax提交的方法
- 07-27Mysql5.7中JSON操作函数使用说明
- 06-28JS中把函数作为另一函数的参数传递方法(总结)
- 06-18帝国CMS7.5版系统模型新增发布后和修改后处理函数扩展
- 04-15php自定义函数实现统计中文字符串长度的方法小结
- 04-03thinkphp 字母函数详解T/I/N/D/M/A/R/U
- 03-15thinkPHP简单调用函数与类库的方法
- 03-10php取得当前时间函数
最近更新
阅读排行
猜你喜欢
- 02-23DedeCMS定时自动生成首页HTML的方法
- 09-16华为mate40pro怎么添加健康码到桌面
- 03-21oppofindx3pro双击锁屏设置步骤教程
- 03-15将 Sublime 打造成一个 Swift 编辑器
- 10-21QQ邮箱怎么设置深色主题
- 01-17同城单身夜约会-同城单身夜约会应用软
- 09-14QQ音乐怎么屏蔽推送消息
- 12-03钉钉会议怎么设置横屏
- 10-19高德地图怎么录制个人语音包
- 07-05Linux fbset命令