剑指 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]

我的解法就比较简单比较粗鲁，虽然能A但是效果就很差

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        
        result = []
        for num in arr:
            if len(result) < k:
                result.append(num)
                continue

            if len(result)>0:
                max_num = max(result)

                if num < max_num:
                    index = result.index(max_num)
                    result[index] = num

        return result

跟着大佬学算法

方法一：排序
本题使用排序算法解决最直观，对数组 arr 执行排序，再返回前 $k$ 个元素即可。使用任意排序算法皆可，本文采用并介绍 快速排序 ，为下文 方法二 做铺垫。

快速排序原理：
快速排序算法有两个核心点，分别为 “哨兵划分” 和 “递归” 。

1. 哨兵划分操作： 以数组某个元素（一般选取首元素）为 基准数 ，将所有小于基准数的元素移动至其左边，大于基准数的元素移动至其右边。
2. 递归： 对 左子数组 和 右子数组 递归执行 哨兵划分，直至子数组长度为 1 时终止递归，即可完成对整个数组的排序。

如下图所示，为示例数组 [2,4,1,0,3,5] 的快速排序流程。观察发现，快速排序和 二分法 的原理类似，都是以 $\log$ 时间复杂度实现搜索区间缩小。

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        def quick_sort(arr, l, r):
            if l >= r:
                return
            
            i, j = l, r
            while i < j:
                while i < j and arr[j] >= arr[l]:
                    j -= 1
                while i < j and arr[i] <= arr[l]:
                    i += 1
                    
                arr[i], arr[j] = arr[j], arr[i]
                
            arr[i], arr[l] = arr[l], arr[i]
            
            quick_sort(arr, l, i - 1)
            quick_sort(arr, i + 1, r)
        
        quick_sort(arr, 0, len(arr) - 1)
        return arr[:k]

注意： 在这里一定要加上等号，这样能保证在排序的列表中如果出现了重复的数字不会在 while i < j 的循环中出不去。另外这两个循环的顺序是不能换的，因为这两个 while执行完， i j 同时指向一个 < arr[l] 的数，因此最后再执行 arr[l], arr[i] = arr[i], arr[l] 可以把哨兵交换到正确的位置。 而如果互换这两句，那么就是 i 先向右遍历，两个 while 执行完， i j 同时指向一个 > arr[l] 的数，那么就不对了。如果要交换写，那么同时也要把哨兵换成数组的末元素，让整个哨兵划分操作对称。

方法二： 基于快速排序的数组划分

题目只要求返回最小的 k 个数，对这 k 个数的顺序并没有要求。因此，只需要将数组划分为 最小的 $k$ 个数 和 其他数字 两部分即可，而快速排序的哨兵划分可完成此目标。

根据快速排序原理，如果某次哨兵划分后 基准数正好是第 $k+1$ 小的数字 ，那么此时基准数左边的所有数字便是题目所求的 最小的 k 个数 。

根据此思路，考虑在每次哨兵划分后，判断基准数在数组中的索引是否等于 $k$ ，若 $true$ 则直接返回此时数组的前 $k$ 个数字即可。

算法流程：
getLeastNumbers() 函数：

1. 若 $k$ 大于数组长度，则直接返回整个数组；
2. 执行并返回 quick_sort() 即可；

quick_sort() 函数：

注意，此时 quick_sort() 的功能不是排序整个数组，而是搜索并返回最小的 $k$ 个数。

1. 哨兵划分：
   • 划分完毕后，基准数为 arr[i] ，左 / 右子数组区间分别为 [l, i - 1][l,i−1] , [i + 1, r][i+1,r] ；
2. 递归或返回：
   • 若 $k<i$ ，代表第 $k+1$ 小的数字在 左子数组 中，则递归左子数组；
   • 若 $k>i$ ，代表第 $k+1$ 小的数字在 右子数组 中，则递归右子数组；
   • 若 $k=i$ ，代表此时 arr[k] 即为第 $k+1$ 小的数字，则直接返回数组前 $k$ 个数字即可；

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        def quick_sort(arr, l, r):
            if l >= r:
                return

            i, j = l, r
            while i < j:
                while i < j and arr[j] >= arr[l]:
                    j -= 1
                    
                while i < j and arr[i] <= arr[l]:
                    i += 1
                
                arr[i], arr[j] = arr[j], arr[i]
                
            arr[i], arr[l] = arr[l], arr[i]
            
            if i > k:
                quick_sort(arr, l, i - 1)
            
            elif i < k:
                quick_sort(arr,  i + 1, r)
            
            else:
                return arr

        
        quick_sort(arr, 0, len(arr) - 1)
        return arr[:k]