获取numpy二维数组每行的N个最大值的索引

直接贴代码了,有问题欢迎留言,一起交流进步。转载请标明出处。

def top_n(rec_list, n):
    '''
    :param rec_list:
    :param n:
    :return:
    '''
    lists = []
    for i in range(len(rec_list)):
        # a = np.argpartition(my_array[i],-3)[-3:]
        lists.append(heapq.nlargest(n, range(len(rec_list[i])), rec_list[i].take))
        # lists[i] = np.append(a)
    return lists

if __name__ =='__main__':
    my_array = np.array([[2, 5, 1, 3, 9], [6, 2, 0, 12, 8], [1, 6, 7, 11, 2]])
    lists = top_n(my_array, 3)
    print(lists)

实验结果如下所示,此时的lists的数据类型是列表。
这是运行的结果

你可能感兴趣的:(python,python)