python获取多维数组中每一行最大的前n个数以及其对应的下标

业务背景

最近在做一个图像识别模型的业务，具体的业务场景为给出一张用餐托盘图片，要识别出托盘上每道菜品。领导要求要得到每道菜品前10个识别的结果，按概率从大到小排列。

技术分析

识别后最终的结果都被封装到output里面。每一列表示每一道菜，每一行表示每一道菜的识别概率，行的下标对应着录入菜品种类。于是乎就把问题转变为遍历N维数组，得到每一行的前十个最大值以及对应的下标。

output = nn.Softmax(dim=1)(output).cpu().detach().numpy()

我网上扒拉了一下，有一个解决方案如下：

import heapq

a = [88, 75, 65, 43, 52, 98, 77]

re1 = heapq.nlargest(4, a)  # 求最大的4个元素，并排序
re2 = map(a.index, heapq.nlargest(4, a))  # 求最大的4个索引    nsmallest与nlargest相反，求最小

print(re1)
print(list(re2))  # 因为re2由map()生成的不是list，直接print不出来，添加list()就行了

当时如获至宝，后来真正去用的时候发现了问题，数据类型不对劲儿呐，然后我type(output)了一下，它的数据类型应该是numpy中的一个narray类型的数据，然后用不了。后面看大佬提供的方案，可以转化为list后再使用，成功解决。
具体的转化代码如下：

    for index in output:
        # print(type(index))
        index.tolist()
        # print(index)
        re1 = heapq.nlargest(10, index.tolist())  
        re2 = map(index.tolist().index, heapq.nlargest(10, index.tolist()))  
        print(re1)
        print(list(re2))  

这样就得到了我们想要的数据了。至于后面把数据封装一下，再返回给前端那就是小问题了。