剑指 Offer 40. 最小的k个数

题目描述

https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof
输入整数数组 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个数外，还可以使用基于快排的分治算法，使得左边为最小的四个数字的集合。
先复习一下快速排序：

class Solution:
    def Partition(self, arr, low, high):
        pivot = arr[low]
        while low < high:
            # 啥时候需要换?high指向的太小了,比pivot都小就需要换,比pivot大或相等就不用换
            while low < high and pivot <= arr[high]:
                high -= 1
            arr[low], arr[high] = arr[high], arr[low]
            # 啥时候需要换?low指向的太大了,比pivot都大就需要换,比pivot小或相等就不用换
            while low < high and pivot >= arr[low]:
                low += 1
            arr[low], arr[high] = arr[high], arr[low]
        return low

    def QuickSort(self, arr: list[int], low: int, high: int):
        if low < high: # 注意这里有个if
            pivot_pos = self.Partition(arr, low, high)
            self.QuickSort(arr, low, pivot_pos - 1) # 注意这里递归的不是Partition
            self.QuickSort(arr, pivot_pos + 1, high)

在QuickSort中添加一句判断：

# 确保前k个数最小就停止
if pivot_pos == k:
     return

即可确保在本次递归中，递归的数组中从low到k是比k到high小的。（注意，不一定是一次满足if就可以做到0到k比k到右边小，因此还是要反复递归）。
在所有递归完成后，返回数组前k个数就可以了。

解题代码

class Solution:
    def Partition(self, arr, low, high):
        pivot = arr[low]
        while low < high:
            # 啥时候需要换?high指向的太小了,比pivot都小就需要换,比pivot大或相等就不用换
            while low < high and pivot <= arr[high]:
                high -= 1
            arr[low], arr[high] = arr[high], arr[low]
            # 啥时候需要换?low指向的太大了,比pivot都大就需要换,比pivot小或相等就不用换
            while low < high and pivot >= arr[low]:
                low += 1
            arr[low], arr[high] = arr[high], arr[low]
        return low

    def QuickSort(self, arr: list[int], low: int, high: int, k: int):
        if low < high:
            pivot_pos = self.Partition(arr, low, high)
            # 确保前k个数最小就停止
            if pivot_pos == k:
                return
            else:
                self.QuickSort(arr, low, pivot_pos - 1, k)
                self.QuickSort(arr, pivot_pos + 1, high, k)

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

执行结果

执行结果：通过
执行用时：300 ms, 在所有 Python3 提交中击败了16.10%的用户
内存消耗：15.9 MB, 在所有 Python3 提交中击败了44.23%的用户