题目描述
输入整数数组 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]
限制:
0 <= k <= arr.length <= 10000
0 <= arr[i] <= 10000
解答方法
方法一:快排
思路
代码
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
def partition(arr,l, r):
pivot = arr[r]
i = l-1
for j in range(l,r):
if arr[j]< pivot:
i += 1
arr[i], arr[j] = arr[j], arr[i]
arr [i+1], arr[r] = arr[r], arr[i+1]
return i+1
def random_partition(arr, l, r):
i = random.randint(l,r)
arr[i], arr[r] =arr[r], arr[i]
return partition(arr, l, r)
def quick_sort(arr,l, r, k):
pos = random_partition(arr, l, r)
nums = pos - l + 1
if nums > k:
quick_sort(arr, l, pos-1,k)
elif nums < k:
quick_sort(arr, pos + 1, r, k - nums)
if k==0:
return []
quick_sort(arr, 0, len(arr)-1,k)
return arr[:k]
时间复杂度
期望为 O(n),由于证明过程很繁琐,所以不再这里展开讲。具体证明可以参考《算法导论》第 9 章第 2 小节。
最坏情况下的时间复杂度为 O(n^2)。情况最差时,每次的划分点都是最大值或最小值,一共需要划分 n−1 次,而一次划分需要线性的时间复杂度,所以最坏情况下时间复杂度为 O(n^2)。
空间复杂度
期望为 ,递归调用的期望深度为 ,每层需要的空间为 O(1),只有常数个变量。
最坏情况下的空间复杂度为 O(n)。最坏情况下需要划分 n 次,即 randomized_selected 函数递归调用最深 n - 1 层,而每层由于需要 O(1) 的空间,所以一共需要 O(n)的空间复杂度。
方法二:堆
思路
我们用一个大根堆实时维护数组的前 k 小值。首先将前k 个数插入大根堆中,随后从第 k+1个数开始遍历,如果当前遍历到的数比大根堆的堆顶的数要小,就把堆顶的数弹出,再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可。
Python 语言中的堆为小根堆,因此我们要对数组中所有的数取其相反数,才能使用小根堆维护前 k 小值。
代码
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
if k==0:
return[]
hp = [-x for x in arr[:k]]
#heapify(heap) 让列表具备堆特征
heapq.heapify(hp)
for i in range(k,len(arr)):
if arr[i] < -hp[0]:
#heappop(heap) 从堆中弹出最小的元素
heapq.heappop(hp)
#heappush(heap, x) 将x压入堆中
heapq.heappush(hp, -arr[i])
ans = [-x for x in hp]
return ans
时间复杂度
时间复杂度:,其中 n是数组 arr 的长度。由于大根堆实时维护前 k 小值,所以插入删除都是 的时间复杂度,最坏情况下数组里n 个数都会插入,所以一共需要 的时间复杂度。
空间复杂度
O(k),因为大根堆里最多 k个数。