最小的k个数(剑指offer 40)

Posted 2021-02-20 PangYunsheng

题目描述：输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

方法一：快排

直接通过快排切分排好第 K 小的数（下标为 K-1），那么它左边的数就是比它小的另外 K-1 个数

时间复杂度: 期望O(n)，最坏O(n^2)

空间复杂度: 期望O(logn), 最坏O(n)

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k >= len(arr): return arr
        return self.quick_select(arr, 0, len(arr)-1, k)

    def partition(self, arr, left, right):
        v = arr[left]
        while left < right:
            while left < right and arr[right] >= v: right -= 1
            arr[left] = arr[right]
            while left < right and arr[left] <= v: left += 1
            arr[right] = arr[left]
        arr[left] = v
        return left

    def quick_select(self, arr, left, right, k):
        p = self.partition(arr, left, right)
        if p == k: return arr[:k]
        elif p < k: return self.quick_select(arr, p+1, right, k)
        elif p > k: return self.quick_select(arr, left, p-1, k)

方法二：最大堆

时间复杂度: O(nlogk)

空间复杂度: O(logk)

class Solution:
    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