由刷题学习 heapq

今日一题是 面试题 17.14. 最小K个数 https://leetcode-cn.com/problems/smallest-k-lcci/

还好 提示 0 <= len(arr) <= 100000, 加上可以用一些内置排序, 相对来说还是比较简单的

看到的时候就想到可以使用堆来做, 奈何我只知道但是没用过, 不过这个题目理解真正考察的是快排?

 

做完看官方的解答的时候, 看到了堆的使用方式 

 

class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return list()

        hp = [-x for x in arr[:k]]
        heapq.heapify(hp)
        for i in range(k, len(arr)):
            if -hp[0] > arr[i]:
                heapq.heappop(hp)
                heapq.heappush(hp, -arr[i])
        ans = [-x for x in hp]
        return ans

作者：LeetCode-Solution
链接：https://leetcode-cn.com/problems/smallest-k-lcci/solution/zui-xiao-kge-shu-by-leetcode-solution-o5eg/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

　　

刚开始看的时候不甚理解为啥要取负数(哈哈哈哈哈哈哈, 还是对堆不熟

 

看了下文档就理解了 https://docs.python.org/zh-cn/3/library/heapq.html

堆的特性如下: 保证第0个节点一定是最小的, 但是他不保证最后面的一定是最大的

这样就有趣了, 我们希望拿到的都是最小的, 那么大的就应该排出掉, 怎么让最大的能排在第0个位置呢, 那就是取反

遍历的时候, 如果发现大于新元素, 则将其pop掉, 再新添加一个, 堆还会维护出来最小的(取反后)在第一位, 这样遍历完就可以拿到结果