2021-06-18发表2021-06-18更新刷题7 分钟读完 (大约996个字)

剑指offer-40

剑指 Offer 40. 最小的k个数

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

输入：arr = [0,1,2,1], k = 1
输出：[0]

官方题解

方法一：堆

我们用一个大根堆实时维护数组的前 k 小值。

首先将前 $k$ 个数插入大根堆中，随后从第 $k+1$个数开始遍历，如果当前遍历到的数比大根堆的堆顶的数要小，就把堆顶的数弹出，再插入当前遍历到的数。

最后将大根堆里的数存入数组返回即可。

在下面的代码中，由于 C++ 语言中的堆（即优先队列）为大根堆，我们可以这么做。

而 Python 语言中的堆为小根堆，因此我们要对数组中所有的数取其相反数，才能使用小根堆维护前 $k$ 小值。

堆是一种完全二叉树结构。

堆分为大顶堆和小顶堆，满足Key[i]>=Key[2i+1]&&key>=key[2i+2]称为大顶堆，满足 Key[i]<=key[2i+1]&&Key[i]<=key[2i+2]称为小顶堆。

由上述性质可知大顶堆的堆顶的关键字肯定是所有关键字中最大的，小顶堆的堆顶的关键字是所有关键字中最小的。

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

        hp = [-x for x in arr[:k]]
        heapq.heapify(hp)

        # print(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]
        # print(hp)
        return ans

方法二：快排思想

我们可以借鉴快速排序的思想。

我们知道快排的划分函数每次执行完后都能将数组分成两个部分，小于等于分界值 pivot 的元素的都会被放到数组的左边，大于的都会被放到数组的右边，然后返回分界值的下标。

与快速排序不同的是，快速排序会根据分界值的下标递归处理划分的两侧，而这里我们只处理划分的一边。

我们定义函数 randomized_selected(arr, l, r, k) 表示划分数组 arr 的 [l,r] 部分，使前 k 小的数在数组的左侧，在函数里我们调用快排的划分函数，假设划分函数返回的下标是 pos（表示分界值 pivot 最终在数组中的位置），即 pivot 是数组中第 pos - l + 1 小的数，那么一共会有三种情况：

如果 pos - l + 1 == k，表示 pivot 就是第 k 小的数，直接返回即可；
如果 pos - l + 1 < k，表示第 k 小的数在 pivot 的右侧，因此递归调用 randomized_selected(arr, pos + 1, r, k - (pos - l + 1))；
如果 pos - l + 1 > k，表示第 k 小的数在 pivot 的左侧，递归调用 randomized_selected(arr, l, pos - 1, k)。

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

class Solution:
    def partition(self,nums,l,r):
        pivot = nums[r]
        i = l - 1
        for j in range(l,r):
            # 小于边界值的放到左边
            if nums[j] <= pivot:
                i+= 1 
                nums[i], nums[j] = nums[j], nums[i]
        
        nums[i+1],nums[r] = nums[r], nums[i+1]
        return i + 1
    def randomized_partition(self,nums,l,r):
        i = random.randint(l,r)
        nums[r],nums[i] = nums[i],nums[r]
        return self.partition(nums,l,r)
    
    def randomized_selected(self,arr,l,r,k):
        pos = self.randomized_partition(arr,l,r)
        num = pos-l + 1
        if k < num:
            self.randomized_selected(arr, l, pos - 1, k)
        elif k > num:
            self.randomized_selected(arr, pos + 1, r, k - num)

    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]