numpy高维数组获取top-K

前言

理论知识请自行翻阅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