【剑指 Offer】 40. 最小的k个数

相似题目：

1. ​​【常考排序算法】堆排序​​.
2. ​​【LeetCode】347. 前 K 个高频元素 （全手写）​​.
3. ​​【LeetCode】215. 数组中的第K个最大元素​​.

1.1 直接 sort（）函数

但是面试这样写肯定不行。

补充 sort 与 sorted 区别：

1. sort 是应用在 list 上的方法，sorted 可以对所有可迭代的对象进行排序操作。
2. list 的sort方法返回的是对已经存在的列表进行操作，无返回值，而内建函数sorted方法返回的是一个新的 list，而不是在原来的基础上进行的操作。

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        arr.sort()
        return arr[:k]

1.2 快速排序(排序完再取前 k 个) (时间：)

直接对数组排序，排序后前k个数就是答案，排序一般较快的是，显然这并不是时间复杂度最优解。

1.3 最大堆(排序完再取前 k 个) (时间：)

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        n = len(arr)
        self.buildHeap(arr, n)
        for i in range(n-1, -1, -1):
            self.swap(arr, i, 0)
            self.heapify(arr, i, 0)
        return arr[:k]
    
    def swap(self, arr, i, j):
        arr[i], arr[j] = arr[j], arr[i]
    def buildHeap(self, arr, n):
        for i in range((n-2)//2, -1, -1):
            self.heapify(arr, n, i)
    def heapify(self, arr, n, i):
        c1 = 2 * i + 1
        c2 = 2 * i + 2
        max = i
        if c1 < n and arr[c1] > arr[max]:
            max = c1
        if c2 < n  and arr[c2] > arr[max]:
            max = c2
        if max != i:
            self.swap(arr, max, i)
            self.heapify(arr, n, max)

1.4 最大堆 (时间：; 空间:

【求最大 top k 问题用最小堆，求最小 top k 问题用最大堆】

思路：用一个大顶堆实时维护数组的前

1. 首先 build 一个长度为
2. 随后从第
3. 最后将大根堆里的数存入数组返回即可。（由于 C++ 语言中的堆（即优先队列）为大根堆，我们可以直接操作）。而 Python 语言中库中自带的为小顶堆，因此我们要对数组中所有的数取其相反数，才能使用小顶堆维护前 k 小值。

优点： 适合海量数据求k个最小。因为k个数的堆，空间是固定的，当数组超级大，那么全存入内存都变得不可行的时候，就需要从外存中慢慢读取数字，然后和这个堆进行比较。

而先排序再取前 k 个数的方法就必须吧整个数组放入内存中，才能运行，所以不适合海量数据。

法一：利用函数库 （推荐！！！）：

import heapq
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

法二：手写大顶堆：

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k==0 or len(arr) == 0:
            return []
        n = len(arr)
        nums = arr[:k]          # 需要额外的 O(k) 的空间
        self.buildHeap(nums, k) # 1. 建立长度为 k 的大顶堆
        
        for i in range(k, n):   # 2. 从 k+1 开始遍历，如果小于堆顶元素，则插入堆顶
            if nums[0] > arr[i]:
                nums[0] = arr[i] # 插入堆顶
                self.heapify(nums,k, 0) # 插完以后，恢复大顶堆结构
        return nums              # 3. 返回这个长度为 k 的最大堆
    
    def swap(self, arr, i, j):
        arr[i], arr[j] = arr[j], arr[i]

    def buildHeap(self, arr, k):
        for i in range((k-2) // 2, -1, -1):
            self.heapify(arr, k, i)
    
    def heapify(self, arr, k, i):
        c1 = 2 * i + 1
        c2 = 2 * i + 2
        max = i
        if c1 < k and arr[c1] > arr[max]:
            max = c1
        if c2 < k and arr[c2] > arr[max]:
            max = c2
        if max != i:
            self.swap(arr, max, i)
            self.heapify(arr, k, max)

如果不做

nums = arr[:k]

这一步操作的话，而是直接将 arr[:k] 带入 buildHeap（）函数中，那么buildHeap（）需要返回值，否则，arr[:k] 并没有被改变，具体看以下代码

法三：

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k==0 or len(arr) == 0:
            return []
        n = len(arr)
        # nums = arr[:k]          # 需要额外的 O(k) 的空间
        # self.buildHeap(nums, k) # 1. 建立长度为 k 的大顶堆
        arr[:k] = self.buildHeap(arr[:k], k)
        for i in range(k, n):   # 2. 从 k+1 开始遍历，如果小于堆顶元素，则插入堆顶
            if arr[0] > arr[i]:
                arr[0] = arr[i] # 插入堆顶
                self.heapify(arr,k, 0) # 插完以后，恢复大顶堆结构
        return arr[:k]            # 3. 返回这个长度为 k 的最大堆
    
    def swap(self, arr, i, j):
        arr[i], arr[j] = arr[j], arr[i]

    def buildHeap(self, arr, k):
        for i in range((k-2) // 2, -1, -1):
            self.heapify(arr, k, i)
        return arr
    
    def heapify(self, arr, k, i):
        c1 = 2 * i + 1
        c2 = 2 * i + 2
        max = i
        if c1 < k and arr[c1] > arr[max]:
            max = c1
        if c2 < k and arr[c2] > arr[max]:
            max = c2
        if max != i:
            self.swap(arr, max, i)
            self.heapify(arr, k, max)

上述提到的问题，下面用简单代码说明:

def fun_1(arr):
    arr[0], arr[1] = arr[1], arr[0]

def fun_2(arr):
    arr[0], arr[1] = arr[1], arr[0]
    return arr

a = [1,2,3,4]
fun_1(a[:2])    # a[:2] 没有被真正操作，是生成了一个副本进行操作
print('a:',a)

b = [5,6,7,8]
b[:2] = fun_2(b[:2]) # 函数返回值，再赋给 b[:2]
print('b:',b)

输出：

a: [1, 2, 3, 4]
b: [6, 5, 7, 8]

法四：嵌套函数（推荐）
采用嵌套函数，可以免去考虑上面提到的 arr[:k] 传参过程中的问题

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        def heapify(k, i):
            c1 = 2 * i + 1
            c2 = 2 * i + 2
            max = i
            if c1 < k and arr[c1] > arr[max]:
                max = c1
            if c2 < k and arr[c2] > arr[max]:
                max = c2
            if max != i:
                arr[i], arr[max] = arr[max], arr[i]
                heapify(k, max)

        if k==0 or len(arr) == 0:
            return []
        n = len(arr)
        
        # 建立最大堆
        for i in range((k-2) // 2, -1, -1):
            heapify(k, i)

        for i in range(k, n):   # 2. 从 k+1 开始遍历，如果小于堆顶元素，则插入堆顶
            if arr[0] > arr[i]:
                arr[0] = arr[i] # 插入堆顶
                heapify(k, 0) # 插完以后，恢复大顶堆结构
        return arr[:k]            # 3. 返回这个长度为 k 的最大堆

复杂度：

1. 时间复杂度：

(如果是排序后再选的话，复杂度最低达到 ，当

1.5 类似快速排序的算法 （期望时间：，最坏 ）

class Solution:
    def partition(self, nums, l, r):
        pivot = nums[r]
        i = l - 1
        for j in range(l, r):
            if nums[j] <= pivot:
                i += 1
                nums[i], nums[j] = nums[j], nums[i]
        nums[i + 1], nums[r] = nums[r], nums[i + 1]
        return i + 1

    def randomized_partition(self, nums, l, r):
        i = random.randint(l, r)
        nums[r], nums[i] = nums[i], nums[r]
        return self.partition(nums, l, r)

    def randomized_selected(self, arr, l, r, k):
        pos = self.randomized_partition(arr, l, r)
        num = pos - l + 1
        if k < num:
            self.randomized_selected(arr, l, pos - 1, k)
        elif k > num:
            self.randomized_selected(arr, pos + 1, r, k - num)

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

复杂度分析

1. 时间复杂度：期望为 O(n) ，由于证明过程很繁琐，所以不再这里展开讲。具体证明可以参考《算法导论》第 9 章第 2 小节。
   最坏情况下的时间复杂度为 。情况最差时，每次的划分点都是最大值或最小值，一共需要划分 n−1 次，而一次划分需要线性的时间复杂度，所以最坏情况下时间复杂度为
2. 空间复杂度：O(logn)，递归调用的期望深度为 O(logn)，每层需要的空间为 O(1)，只有常数个变量。
   最坏情况下的空间复杂度为 O(n)。最坏情况下需要划分 n 次，即 randomized_selected 函数递归调用最深 n−1 层，而每层由于需要 O(1) 的空间，所以一共需要 O(n) 的空间复杂度。

参考: ​​LeetCode官方题解​​

大顶堆和这种方法的优劣性比较

在面试中，另一个常常问的问题就是这两种方法有何优劣。看起来分治法的快速选择算法的时间、空间复杂度都优于使用堆的方法，但是要注意到快速选择算法的几点局限性：

第一，算法需要修改原数组，如果原数组不能修改的话，还需要拷贝一份数组，空间复杂度就上去了。

第二，算法需要保存所有的数据。如果把数据看成输入流的话，使用堆的方法是来一个处理一个，不需要保存数据，只需要保存 k 个元素的最大堆。而快速选择的方法需要先保存下来所有的数据，再运行算法。当数据量非常大的时候，甚至内存都放不下的时候，就麻烦了。所以当数据量大的时候还是用基于堆的方法比较好。

参考：

1. ​​【LeetCode】215. 数组中的第K个最大元素​​.
2. ​​【常考排序算法】快速排序​​.
3. ​​【常考排序算法】堆排序​​.

题主本硕机械专业，自学转互联网 算法岗成功，获得阿里、字节、美团、华为等 15+ offer

公众号 「苏学算法」