剑指Offer40.最小的k个数(快排、堆排序)(Python可运行，带测试用例)

一、题目

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

二、题解

1.直接用python自带sort()函数

class Solution:
    def getLeastNumbers(self, arr, k):
        arr.sort()
        # 同arr[0:k]，从开头取到第k个
        return arr[:k]

s = Solution()
print(s.getLeastNumbers([3, 2, 1], 2)) # [1, 2]
print(s.getLeastNumbers([0, 1, 2, 1], 1)) # [0]

2.快速排序

class Solution:
    def quick_sort(self, nums: list, left: int, right: int) -> None:
        if left < right:
            i = left
            j = right
            # 取第一个元素为枢轴量
            pivot = nums[left]
            while i != j:
                # 交替扫描和交换
                # 从右往左找到第一个比枢轴量小的元素，交换位置
                while j > i and nums[j] > pivot:
                    j -= 1
                if j > i:
                    # 如果找到了，进行元素交换
                    nums[i] = nums[j]
                    i += 1
                # 从左往右找到第一个比枢轴量大的元素，交换位置
                while i < j and nums[i] < pivot:
                    i += 1
                if i < j:
                    nums[j] = nums[i]
                    j -= 1
            # 至此完成一趟快速排序，枢轴量的位置已经确定好了，就在i位置上（i和j)值相等
            nums[i] = pivot
            # 以i为枢轴进行子序列元素交换
            self.quick_sort(nums, left, i - 1)
            self.quick_sort(nums, i + 1, right)

    def getLeastNumbers(self, arr, k: int):
        n = len(arr)
        if k <= 0 or arr == []:
            return []
        if n <= k:
            return arr
        # 进行快排
        self.quick_sort(arr, 0, n - 1)
        return arr[:k]

s = Solution()
print(s.getLeastNumbers([3, 2, 1], 2)) # [1, 2]
print(s.getLeastNumbers([0, 1, 2, 1], 1)) # [0]

3.堆排序

import heapq
class Solution:
    def getLeastNumbers(self, arr, k: int):
        if k == 0:
            return list()
        # 求前 K 小，因此用一个容量为 K 的大根堆，每次 pop弹出 出最大的数，那堆中保留的就是前 K 小
        # 使用大根堆来实现，Python语言中的堆为小根堆，因此我们要对数组中所有的数取其相反数，才能使用小根堆维护前 k 小值
        # 将前 k 个数插入大根堆中，随后从第 k+1 个数开始遍历
        hp = [-x for x in arr[:k]]
        # 将列表list类型转化为heap堆, 在线性时间内, 重新排列列表
        heapq.heapify(hp)
        # 从第 k+1 个数开始遍历
        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]
        return ans

s = Solution()
print(s.getLeastNumbers([3, 2, 1], 2)) # [2, 1]
print(s.getLeastNumbers([0, 1, 2, 1], 1)) # [0]

时间复杂度：O(nlogk)，其中 n 是数组 arr 的长度。由于大根堆实时维护前 k 小值，所以插入删除都是 O(logk) 的时间复杂度，最坏情况下数组里 n 个数都会插入，所以一共需要 O(nlogk) 的时间复杂度。

空间复杂度：O(k)，因为大根堆里最多 k 个数。