前言
理论知识请自行翻阅numpy的argpartition和partition方法的实现原理,该文章仅仅包含使用和效率验证。此外,numpy版本需要>=1.8.0。
正文
不废话了,直接放代码,一看就懂,看不懂再说,自己跑一下就知道。
import numpy as np
def get_sorted_top_k(array, top_k=1, axis=-1, reverse=False):
"""
多维数组排序
Args:
array: 多维数组
top_k: 取数
axis: 轴维度
reverse: 是否倒序
Returns:
top_sorted_scores: 值
top_sorted_indexes: 位置
"""
if reverse:
# argpartition分区排序,在给定轴上找到最小的值对应的idx,partition同理找对应的值
# kth表示在前的较小值的个数,带来的问题是排序后的结果两个分区间是仍然是无序的
# kth绝对值越小,分区排序效果越明显
axis_length = array.shape[axis]
partition_index = np.take