输入整数数组 arr ,找出其中最小的 k 个数。例如,输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。
思路1:先进行排序
使用sort()函数
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
arr.sort()
res=[]
for i in range(k):
res.append(arr[i])
return res
不使用sort()函数
快速排序算法见我的另一篇python快速排序实现
思路2
堆排序(适合海量数据)
由于大根堆实时维护前 k 小值,所以插入删除都是 O(logk) 的时间复杂度,最坏情况下数组里 n 个数都会插入,所以一共需要 O(nlogk) 的时间复杂度。
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
if k == 0:
return list()
hp = [-x for x in arr[:k]]
heapq.heapify(hp)
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