[LeetCode]（面试题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

解题思路

解法一：排序

对原数组从小到大排序后取出前 k 个数即可。
复杂度分析：
时间复杂度：O(nlogn)，其中 n 是数组 arr 的长度。算法的时间复杂度即排序的时间复杂度。
空间复杂度：O(logn)，排序所需额外的空间复杂度为 O(logn)。

解法二：堆排序

堆的定义：
堆通常是一个可以被看做一棵树的数组对象。堆总是满足下列性质：
1）堆中某个节点的值总是不大于或不小于其父节点的值，即堆可分为最大堆（大根堆、大顶堆）或者最小堆（小根堆、小顶堆）
2）堆总是一棵完全二叉树。
因为堆是一棵完全二叉树，所以一般可以用数组来实现。数组的下标对应堆中节点的编号。为方便起见，我们假设数组下标从 1 开始。那么对于堆中每个节点与其左右子节点的编号关系都有：

1. leftID = fatherID * 2
2. rightID = fatherID * 2 + 1
3. fatherID = sonID / 2

有了数据存储格式之后，定义以下堆方法：

1. int size() { … }; 返回堆内元素个数。
2. int top() { … }; 返回根节点的元素。
3. void push(int x) { … }; 插入一个元素。
4. void pop(int x) { … }; 将根节点元素从堆中弹出。

前两种方法较简单，size() 可以维护一个计数，在 push 和 pop 时更新即可。top() 直接返回根节点的元素即可。
主要讲下第三和第四个灵魂函数：

push 方法
由于性质二的限制，push后堆也应该是一棵完全二叉树，所以必须将元素追加到数组末尾。
又由于性质一的限制，需要对刚刚push的元素进行冒泡。
以最大堆为例，设刚刚push的元素的编号为 id，val[id] 表示对应节点的值：
1）如果 id == 1，数组只有一个元素，冒泡过程结束。
2）如果 val[id] > val[id/2，那么需进行交换，swap(val[id], val[id/2])，id /= 2，跳转第 1 步；否则，当前元素值不大于父节点的值，满足性质一，算法结束。

pop 方法
pop 需要分两步走：
第一步，先将根节点与编号最大节点的元素互换，并删除编号最大的节点。
此时堆仍然是一棵完全二叉树，但有可能不满足性质一。
所以我们需要对根节点的元素进行下沉操作，以大顶堆为例，设置一个游标 id, 初始指向根节点：
1）如果id指向叶子节点，算法结束。
2）如果id指向节点大于其左右子节点的值，即已经满足性质一了，算法结束。
3）设id的左右子节点中，拥有较大值的编号为 p，交换 id 与 p 的值，并将 id 指向 p 节点。跳转步骤 1。

topk 问题一般用堆可解。求最小的 k 个元素可以使用大顶堆解决，反之求最大的 k 个元素，可用小顶堆解决。
以本题为例，我们可以遍历数组 arr，对其元素执行 push 操作。每次push后，检查size，若 size > k，则执行 pop 操作。由于每次从堆顶弹出的数都是堆中最大的，最小的 k 个元素一定会留在堆里。这样，把数组中的元素全部入堆之后，堆中剩下的 k 个元素就是最小的 k 个数了。（在代码上也可以做一些优化，如果当前数字不小于堆顶元素，数字可以直接丢掉，不入堆。）
由于 Python 语言中的堆为小根堆，因此我们要对数组中所有的数取其相反数，才能使用小根堆维护前 k 小值。

复杂度分析：
时间复杂度：入堆和出堆操作的时间复杂度均为 O(logk)，每个元素都需要进行一次入堆操作，故算法的时间复杂度为 O(nlogk)。
空间复杂度：由于使用了一个大小为 k 的堆，空间复杂度为 O(k)。

解法三：快速选择

“查找第 k 大的元素”是一类算法问题，称为选择问题。找第 k 大的数，或者找前 k 大的数，有一个经典的 quick select（快速选择）算法。这个名字和 quick sort（快速排序）看起来类似，也都是分治法的思想。
在快速排序中有一步很重要的操作是 partition（划分），从数组中随机选取一个枢纽元素 v，然后原地移动数组中的元素，使得小于等于v 的元素在 v 的左边，大于v 的元素在 v 的右边。这个 partition 操作是原地进行的，需要 O(n) 的时间，接下来，快速排序会递归地排序左右两侧的数组。而快速选择（quick select）算法的不同之处在于，接下来只需要递归地选择一侧的数组。快速选择算法想当于一个“不完全”的快速排序，因为我们只需要知道最小的 k 个数是哪些，并不需要知道它们的顺序。
定义函数 randomized_selected(arr, p, q, k) 表示划分数组 arr 的 [p,q] 部分，使前 k 小的数在数组的左侧，在函数里我们调用快排的划分函数，假设划分函数返回的下标是 m（表示分界值 pivot 最终在数组中的位置），即 pivot 是数组中第 m - p + 1 小的数，那么一共会有三种情况：
1）如果 m - p + 1 == k，表示 pivot 就是第 k 小的数，直接返回即可；
2）如果 m - p + 1 > k，表示第 k 小的数在 pivot 的左侧，递归调用 randomized_selected(arr, p, m - 1, k)；
3）如果 m - p + 1 < k，表示第 k 小的数在 pivot 的右侧，因此递归调用randomized_selected(arr, m + 1, q, k - (m - p + 1))。

函数递归入口为 randomized_selected(arr, 0, arr.length - 1, k)。在函数返回后，将前 k 个数放入答案数组返回即可。

复杂度分析：
时间复杂度：期望为 O(n)。最坏情况下的时间复杂度为 O(n^2)。情况最差时，每次的划分点都是最大值或最小值，一共需要划分 n - 1 次，而一次划分需要线性的时间复杂度，所以最坏情况下时间复杂度为 O(n^2)。
空间复杂度：期望为 O(logn)，递归调用的期望深度为 O(logn)，每层需要的空间为 O(1)，只有常数个变量。最坏情况下的空间复杂度为 O(n)。最坏情况下需要划分 n 次，即 randomized_selected 函数递归调用最深 n−1 层，而每层由于需要 O(1) 的空间，所以一共需要 O(n) 的空间复杂度。

注：堆排序与快速选择方法的优劣性比较

在面试中，另一个常问的问题就是这两种方法有何优劣。看起来分治法的快速选择算法的时间、空间复杂度都优于使用堆的方法，但是要注意到快速选择算法的几点局限性：
第一，算法需要修改原数组，如果原数组不能修改的话，还需要拷贝一份数组，空间复杂度就上去了。
第二，算法需要保存所有的数据。如果把数据看成输入流的话，使用堆的方法是来一个处理一个，不需要保存数据，只需要保存 k 个元素的最大堆。而快速选择的方法需要先保存下来所有的数据，再运行算法。当数据量非常大的时候，甚至内存都放不下的时候，就麻烦了。所以当数据量大的时候还是用基于堆的方法比较好。

代码

解法一：排序

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]:
        if k==0:
            return list()
        # 取相反数创建小根堆，在小根堆中的是最大的k个数，即原数组的最小的k个数
        h = [-arr[i] for i in range(k)]
        heapq.heapify(h)
        for j in range(k,len(arr)):
            if -h[0]>arr[j]:
                heapq.heappop(h)
                heapq.heappush(h,-arr[j])
        res = [-x for x in h]
        return res

解法三：快速选择

class Solution:
    # 划分函数，返回基准在数组中的下标
    def partition(self, nums, p, q):
        pivot = nums[p]
        while p!=q:
            while p<q and nums[q]>pivot:
                q -= 1
            nums[p] = nums[q]
            while p<q and nums[p]<=pivot:
                p += 1
            nums[q] = nums[p]
        nums[p] = pivot
        return p
    
    # 划分数组，使前k小的数全在数组左侧
    def randomized_selected(self, arr, p, q, k):
        m = self.partition(arr, p, q)
        # 第k小的数在数组左侧，在左侧递归查找
        if (m-p+1)>k:
            self.randomized_selected(arr, p, m-1, k)
        # 第k小的数在数组右侧，在右侧递归查找
        elif (m-p+1)<k:
            self.randomized_selected(arr, m+1, q, k-(m-p+1))

    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k==0:
            return list()
        self.randomized_selected(arr, 0, len(arr)-1, k)
        return arr[:k]