题目描述
输入整数数组 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的元素 小的元素都放左边, 大的都放右边, 这样左边部分的元素就是最终结果. 根据快排的思想操作过程如下:
- 把数组第一个元素当做哨兵, 进行第一次快排, 将所有比哨兵小的元素放到左边, 比哨兵大的元素放在右边, 最后将哨兵与i位置交换
- 如果k比i小, 说明我们需要的哨兵位置在左边部分, 进行左边部分的递归
- 如果k比i大, 说明我们需要的哨兵位置在右边部分, 进行右边部分的递归
- 如果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)