输入整数数组 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] class Solution: def getLeastNumbers(self, arr: List[int], k: int) -> List[int]: '''堆排序''' res = [] # 要返回的list '''这三步是将数组整体后移一个位置,第一个数存在下标为1的位置''' arr.append(0) for i in range(len(arr)-2, -1, -1): arr[i+1] = arr[i] # 建立大根堆 def BuildMaxHeap(array: List[int], length): for i in range(len(array)//2 - 1, 0, -1): # 从i=[length/2]~1反复调整堆 HeadAdjust(array, i, length) def HeadAdjust(array, j, length): '''将元素j为根的子树进行调整''' array[0] = array[j] # 先用array[0]暂存要调整的节点的值 i = 2*j while i < length: if i < length - 1 and array[i]>array[i+1]: # 取较小值的下标 i += 1 if array[0] <= array[i]: # 如果要调整的值本就是最小,则返回 break else: array[j] = array[i] # 否则,修改j的值,继续向下调整 j = i i *= 2 array[j] = array[0] # 被调整的节点放入最终位置 BuildMaxHeap(arr, len(arr)) # 先建立小根堆 for i in range(len(arr)-1, len(arr)-k-1, -1): # 找出前k个最小的值 res.append(arr[1]) arr[1], arr[i] = arr[i], arr[1] # 交换 HeadAdjust(arr, 1, i) # 再调整堆 return res