剑指 Offer 40. 最小的k个数 - python3 堆解法

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        import heapq
        # 由于python的heapq堆维护的是一个从小到大的优先队列，
        # 而题目需要保留最小的k个数，因此将所有数取反
        heap = [-x for x in arr[:k]]
        heapq.heapify(heap)
        for i in arr[k:]:
            # METHOD:1 先入队列，后弹出队首元素
            # heapq.heappush(heap, -i) # 将元素取反放入队列，并维护优先队列的有序性
            # heapq.heappop(heap) # 弹出最前面的数

            # METHOD:2 入队列弹出队首元素，同时进行，提高性能！
            heapq.heappushpop(heap, -i)
        return [-x for x in heap]

【 注 】看到官方题解中有判断是否入队列：

if -i > heap[0]:

尝试加上了判断是否入队列后并没有变快，查看python的源代码发现是heapq.heappushpop方法中已经包含了这个逻辑：

def heappushpop(heap, item):
    """Fast version of a heappush followed by a heappop."""
    if heap and heap[0] < item:
        item, heap[0] = heap[0], item
        _siftup(heap, 0)
    return item

Ref: https://github.com/python/cpython/blob/3.9/Lib/heapq.py

 

原题目

剑指 Offer 40. 最小的k个数

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

 

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

 

限制：

• 0 <= k <= arr.length <= 10000
• 0 <= arr[i] <= 10000