python list Topk用法学习笔记

python topk 排序

对于一个python list 或者numpy数组，我需要找到这个list中最大的K个数及其对应的下标。

解决方式：
1. 可以构造字典通过排序解决，不过代码量较多。

2. 使用heapq库，可以直接获取最大值的下标和数值。

如果要取最小的数，使用 nsmallest即可。

1维数据topk：

import heapq
a = [4 ,2 ,6 ,1 ,9 ,9]

# 获取下标， 输出为[4, 5, 2]
top_index= heapq.nlargest(3, range(len(a)), a.__getitem__)

print(top_index)
# 获取数值， 输出为[9, 9, 6]
result=heapq.nlargest(3 ,a)

print(result)

多维也是ok的，默认第1列排序：

if __name__ == '__main__':
    import heapq

    a = [[8,'2'], [3,'5'], [3,'5'], [5,'6'], [9,'4'], [10,'1']]

    # 获取下标， 输出为[5, 4, 0]
    aaa=heapq.nlargest(3, range(len(a)), a.__getitem__)

    # 获取数值， 输出为[9, 9, 6]
    bbb=heapq.nlargest(3, a)

    print(aaa)
    print(bbb)

每列的数据类型需要一样，否则会报错：

if __name__ == '__main__':
    import heapq

    a = [[8,'2'], [3,'5'], [3,5], [5,6], [9,4], [10,1]]

    # 获取下标， 输出为[5, 4, 0]
    aaa=heapq.nlargest(3, range(len(a)), a.__getitem__)

    # 获取数值， 输出为[9, 9, 6]
    bbb=heapq.nlargest(3, a)

    print(aaa)
    print(bbb)

字符串类型：也可以，只是排序依据有待确认，默认第1列排序

if __name__ == '__main__':
    import heapq

    a = [['8','2'], ['3','5'], ['3','5'], ['5','6'], ['9','4'], ['10','1']]

    # 获取下标， 输出为[5, 4, 0]
    aaa=heapq.nlargest(3, range(len(a)), a.__getitem__)

    # 获取数值， 输出为[9, 9, 6]
    bbb=heapq.nlargest(3, a)

    print(aaa)
    print(bbb)

第2维排序方法：

if __name__ == '__main__':

    import numpy as np
    import heapq

    a = [['8','2'], ['3','8'], ['3','5'], ['5','6'], ['9','4'], ['10','1']]

    a=np.array(a)[:, 1]

    # 获取下标， 输出为[5, 4, 0]
    aaa=heapq.nlargest(3, range(len(a)), a.__getitem__)

    # 获取数值， 输出为[9, 9, 6]
    bbb=heapq.nlargest(3, a)

    print(aaa)
    print(bbb)