对数据进行排序是一个很常见的需求,但有时候我们并不需要对完整的数据进行排序,只需要排前几的数据,也就是经典的 Top-K 问题。
Top-K 问题的经典解法有两种:一种是脱胎于快速排序(Quick Sort)的快速选择(Quick Select)算法,核心思路是在每一次Partion
操作后下一次递归只操作前K
项数据。另一种是基于堆排序的方法。
Python 中有两个标准库可以原生的支持堆排序(优先队列),分别是heapq
和PriorityQueue(queue)
。
heapq
heapq
标准库提供了一些工具函数用来对list
对象实现二叉堆的各种操作(就地修改list
对象)。简单的用法如下:
建堆
1 | import heapq |
获取堆顶元素
1 | assert heapq.heappop(arr) == 0 |