面试题40 最小的k个数 堆排序/ sort函数

输入整数数组 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]

思路1 排序+取前k个数

根据题目要求，先将arr按从小到大排序，再取前k个数，即可得到所求结果。这应该是最显然的思路，并且用python很好实现。

时间复杂度同样是O(Nlog2N)

空间复杂度是O(log2N)~O(N)

代码

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        arr.sort()
        return arr[:k]

上述代码也可以简化为

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        return sorted(arr)[:k]

结果

 

思路2 大顶堆

堆是一种非线性结构，（本篇随笔主要分析堆的数组实现）可以把堆看作一个数组，也可以被看作一个完全二叉树，通俗来讲堆其实就是利用完全二叉树的结构来维护的一维数组

按照堆的特点可以把堆分为大顶堆和小顶堆

大顶堆：每个结点的值都大于或等于其左右孩子结点的值

小顶堆：每个结点的值都小于或等于其左右孩子结点的值

通过构造具有k个元素的大顶堆来返回题目要求的最小k个元素。具体步骤如下：

取数组前k个元素初始化堆，从最后一个非叶结点开始到根结点来构建最大堆；
当某个元素大于堆顶元素时，直接抛弃；
当某个元素小于堆顶元素时，替换堆顶元素，再从堆顶重新构建最大堆；

可以python标准库heapq.nsmallest(k, arr)，一行代码就可以解决
之前不熟悉堆，明天再来写一下堆排序算法。