leetcode剑指 Offer 40. 最小的k个数

题目要求

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

解题思路

方法一

排序后，返回最小的k个数。

class Solution(object):
    def getLeastNumbers(self, arr, k):
        """
        :type arr: List[int]
        :type k: int
        :rtype: List[int]
        """
        arr.sort()
        return arr[:k]

方法二

1. 使用堆排序。
2. python中heapq库实现堆，但只能实现小顶堆。需要对原数组进行取负号操作。

class Solution(object):
    def getLeastNumbers(self, arr, k):
        """
        :type arr: List[int]
        :type k: int
        :rtype: List[int]
        """
        # k为0时，特殊处理。
        if not k:
            return []
        # 导入heapq库
        import heapq
        # 由于需要heapq只能实现小顶堆，需取负号。
        for i in range(len(arr)):
            arr[i] = -arr[i]
        res = arr[:k]
        # 先将前k个数，建立堆。
        heapq.heapify(res)
        # 在小顶堆中，始终保存最大的k个数字。
        for i in range(k, len(arr)):
            if arr[i] > res[0]:
                heapq.heappop(res)
                heapq.heappush(res, arr[i])
        out = []
        # 最后输出堆中数字的相反数，为最小的k个数。
        for x in res:
            out.append(-x)
        return out

方法三

利用快速排序的变形方法

class Solution(object):
    def getLeastNumbers(self, arr, k):
        """
        :type arr: List[int]
        :type k: int
        :rtype: List[int]
        """
        # 排序
        self.random_select_sort(arr, 0, len(arr)-1, k)
        # 返回前k项
        return arr[:k]
    
    def random_select_sort(self, nums, l, r, k):
        # 一定要考虑下标问题
        if l >= r:
            return
        # 分治处理，并获取中间元素，其左边的元素小，右边的元素大于等于中间元素。
        mid = self.recursive_sort(nums, l, r)
        # 注意次数mid-l的处理，左侧k个以上，则缩小排序范围。
        if mid-l > k-1:
            self.random_select_sort(nums, l, mid-1, k)
        # 左侧不足k个元素，则左侧的无需再排序，只需在右侧排序。寻找k-(mid-l+1)。
        elif mid-l < k-1:
            self.random_select_sort(nums, mid+1, r, k-(mid-l+1))

    def recursive_sort(self, nums, l, r):
        # 指针i初始赋值l-1。
        i = l - 1 
        # 遍历所有元素。
        for j in range(l, r):
            if nums[j] < nums[r]:
                i += 1
                # 交换元素位置。
                nums[j], nums[i] = nums[i], nums[j]
        # 最后将最后的值，即默认为pivot值的元素放到正确位置，使得其左边小于，右边不小于。
        i += 1
        nums[i], nums[r] = nums[r], nums[i]
        return i