Python 查找 List 中 TopK 数值与索引的高效方法
本文介绍在 Python 中查找列表或数组前 K 个最大/最小值及其索引的两种方案。首先演示利用 heapq 库直接获取数值与下标的高效用法，随后深入解析基于中间值划分的 TopK 算法原理与递归实现，适用于大规模数据场景下的性能优化。
需求：
针对 Python 列表或 NumPy 数组，目标是提取其中最大的 K 个数值，并同步获取它们对应的原始索引位置。
解决方式：
1. 可通过构建字典并进行排序来实现，但代码相对繁琐。
2. 推荐直接使用 heapq 库，该库能高效地直接返回最大值的索引和具体数值。
import heapq
a = [4,2,6,1,9,9]
# 获取下标， 输出为[4, 5, 2]
heapq.nlargest(3, range(len(a)), a.__getitem__)
# 获取数值， 输出为[9, 9, 6]
heapq.nlargest(3,a)若需获取最小的数，调用 nsmallest 函数即可。
补充：Python 利用中间值求 TopK 算法
算法思想
首先需明确核心目标：解决什么问题？
TopK 问题的核心在于从一组数据中筛选出前 K 个最大值或最小值。在此过程中，需考虑数据是否包含重复元素，以及是否需要进行去重处理。
明确了处理目标后，接下来探讨 Python 中处理 TopK 问题的具体算法流程。
若直接采用全排序，则无需专门引入 TopK 算法。当数据量庞大时，仅选取 TopK 可大幅减少排序计算量。其优化程度取决于具体场景，例如组合 C 与排列 A 的区别，前者仅抽取元素，后者则涉及抽取后的排列顺序。
以下以查找 TopK 最大值为例进行说明，查找最小值的逻辑只需稍作调整即可。
该算法的核心思想是利用中间值将数列划分为三个部分：
【比中间值大的列表】、中间值本身、【比中间值小的列表】。
当进行如下比较时：
【比中间值大的列表】的元素个数等于 k
时，即可确定前 K 个最大值。因此，
算法的关键在于准确找出这个中间值。
如何找出中间值
以列表的第一个元素作为初始中间值，将其拆分为上述三个部分。
if 【比中间值大的列表】的元素个数 == k：return 中间值 #程序出口，结束递归。
if 【比中间值大的列表】的元素个数 < k ：
·····继续在【比中间值小的列表】中查找
·····需要查找的数量为 K - 【比中间值大的列表】的元素个数 - 1
（为何要减一？因为前一次选取的中间值本身不属于“比中间值大”或“比中间值小”的列表，即它被排除在前后两部分之外，因此需扣除该中间值的位置。）
if 【比中间值大的列表】的元素个数 > k ：
……也就是说，若比中间值大的列表元素数量超过 K，则只需在该列表中继续查找即可。
结合代码和注释进行理解。
若要查找最小值，只需对逻辑稍作修改即可。此外，还可引入布尔参数来指定是查找前 K 个最大值还是最小值。
#2019 11 04
#author 半斤地瓜烧
#TopK 算法，找出序列中前K个最大值的
#输入一个seq
# 输出以seq[0]为中间值 划分的三个部分，中间值，比这个值大的seq ,比这个值小的seq，
# 即splitNum,theBig,theSmall
def Split_Seq(seq):
splitNum = seq[0]
seq = seq[1:]#两个部分都不包含中间值，因此切片去除seq[0]
theBig = [x for x in seq if x >= splitNum]
theSmall = [x for x in seq if x < splitNum]
return splitNum,theBig,theSmall
#找出中间值
def topKNum(seq,k):
splitNum, theBig, theSmall = Split_Seq(seq)
theBigLen = len(theBig)
if k == theBigLen:
return splitNum#出口,返回这个中间值,
# 为什么不直接返回thebig?因为存在递归的原因thebig 不是在初始的seq找出来的
#需要重新Split,即可，读者自己思考
# 大值的列表中还未够K个数的情况，
if k > theBigLen:
return topKNum(theSmall,k-theBigLen-1)
# 大值的列表中大于K个数的情况
return topKNum(theBig,k)
#由中间值找出TopK个值，<list>
def getTopK(seq,k):
return [i for i in seq if i > topKNum(seq, k)]
if __name__ == '__main__':
alist = [7, 3, 5, 1,885,234,2211,222,22, 2, 11, 2, 115]
print("===为了验证，引入排序观看===", sorted(alist,reverse= True))
print(getTopK(alist, 3))以上为个人经验总结，希望能为大家提供参考，也希望大家多多支持本站。
免责声明：本文内容仅供参考，不代表本站观点。转载或引用请注明出处。
