剑指 Offer 40. 最小的k个数

300 阅读2分钟

题目描述

输入整数数组 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]

解题思路1: 排序法 O(nlogn)

将数组排序, 直接返回数组的前k个元素

示例代码1:

def getLeastNumbers(self, arr: [int], k: int) -> [int]:
    arr.sort()
    return arr[:k]

解题思路2: 最大堆 O(nlogk)

构造一个有k个元素最大堆, 然后遍历整个数组, 如果堆中元素不超过k个, 就一直加入; 超过以后, 如果当前元素大于堆顶元素, 则直接跳过, 如果小于堆顶元素, 则将删除堆顶元素, 将当前元素入堆, 并且重新调整最大堆. 遍历结束后, 堆里的元素就是最小的k个元素

示例代码2: python的堆结构直接提供了返回堆的最小k个数, 可以直接使用

def getLeastNumbers(self, arr: [int], k: int) -> [int]:
    return heapq.nsmallest(k, arr)

解题思路3: 类快排思想

因为返回结果不要求顺序, 所以我们可以结合快排的思想, 将比 排序后下标为k的元素 小的元素都放左边, 大的都放右边, 这样左边部分的元素就是最终结果. 根据快排的思想操作过程如下:

  1. 把数组第一个元素当做哨兵, 进行第一次快排, 将所有比哨兵小的元素放到左边, 比哨兵大的元素放在右边, 最后将哨兵与i位置交换
  2. 如果k比i小, 说明我们需要的哨兵位置在左边部分, 进行左边部分的递归
  3. 如果k比i大, 说明我们需要的哨兵位置在右边部分, 进行右边部分的递归
  4. 如果k与i相等, 那么说明我们就找到了最终的哨兵位置, 返回原数组k之前的元素

示例代码3:

def getLeastNumbers(self, arr: [int], k: int) -> [int]:
    if k >= len(arr):
        return arr

    def quick_sort(l, r):
        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[l], arr[i] = arr[i], arr[l]

        if k < i:
            quick_sort(l, i - 1)
        if k > i:
            quick_sort(i + 1, r)
        return arr[:k]

    return quick_sort(0, len(arr) - 1)