获取list中Top K个值对应的索引

背景

在得到概率分布的list之后，想要得到 Top K个概率值及其索引，因为索引对应着label id。

方法

    import numpy as np
    p = [0.07, 0.2, 0.1, 0.03, 0.1, 0.5]
    topk = 2
    print("方法1：")
    index_list = sorted(range(len(p)), key=lambda i: p[i])[-topk:] # sorted(range(len(p)), key=lambda i: p[i], reverse=True)[:2]
    pro_list = np.array(p)[index_list]
    print("indexs={}".format(index_list))
    print("pro values={}".format(pro_list))

    import numpy as np
    ind = np.argpartition(p, -topk)[-topk:] # 返回结果是 array
    print("方法2：")
    print("indexs={}".format(ind.tolist()))
    # preds_ind = p[ind]
    # print(preds_ind)

    # 方法3：
    print("方法3：")
    inds = np.argsort(p)[-topk:]
    print("indexs={}".format(inds.tolist()))

运行结果：

方法1：
indexs=[1, 5]
pro values=[0.2 0.5]
方法2：
indexs=[1, 5]
方法3：
indexs=[1, 5]