堆排序——最小的k个数（python）

一. 题目描述

输入整数数组 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

二. 解题思路

1. 最初的想法（小顶堆）

首先自然而然联想到维护k个节点的一个小顶堆，创建完之后再去维护k个节点的小顶堆，维护完之后这k个节点就是我们要找的那k个数字。
但是有个问题，以4,5,1,6,2,7,3,8为例，要求求最小的4个数字，所以先选数组中的前4个数字去组成四个数字的小顶堆，然后去遍历后面的n-k个数，去维护这个小顶堆，我们判断当前遍历到的节点小于小顶堆的堆顶时，那肯定是要放进来维护的，但是维护起来很麻烦；但是若大于小顶堆的堆顶，就不确定是否能放进来，因为只知道大于最小的数，但是不确定是否在这k个数中，所以存在一个选择的困难。

2. 换方向（大顶堆）

所以我们选择大顶堆维护，还是这两种情况：遍历剩余的n-k个元素时，若当前元素大于大顶堆的堆顶元素，那么当前这个数字肯定不属于最小的k个数字，就不放入堆里面；若小于，说明当前位置的元素是有可能属于最小的k个元素的。

问题来了：

（1）怎么样去维护一个堆，让它成为大顶堆？

（若root节点的编号为0，那么对于每一个节点而言，他的左孩子节点编号时2 * index+1，右孩子节点编号是2*index+2）

【非叶子节点下标=节点个数//2-1 ：即若root从0开始，则 从 0~节点个数//2-1 这个闭区间范围内都是非叶子节点，节点个数//2-1之后就是叶子节点了】

大顶堆的维护：自底向上的维护，对于叶子节点而言没有左孩子右孩子，因为大顶堆要求左孩子右孩子都小于父节点，所以不考虑叶子节点，直接从非叶子节点开始。
大顶堆的建立总结：从非叶子节点开始维护，维护的过程中根据大顶堆的性质（节点元素大于左右节点元素的值）判断当前节点应该处于大顶堆的什么位置。

（2）具体实现：

遍历n-k个数字时，若当前数字大于大顶堆的堆顶，则不添加进堆中，跳过；若小于，则直接将堆顶替换成当前数字再进行后续操作（即只是替换了根节点，他的左孩子和右孩子还是满足大顶堆的性质的），后续操作即比较新的根节点和左右节点的大小，看是否需要更换位置（本质就是维护index=0的情况）

• 先判断特殊情况，当k大于输入数组的长度，或者k=0，这时直接返回[]即可，比如：数组长度为5，让你返回最小的6个数，自然是找不到的，就返回[]即可。
• 然后用数组去模拟k个节点的堆结构，先从原数组中取出前k元素（可以遍历也可以用切片索引复制）；下面就开始维护堆使其成为大顶堆（我们是从非叶子节点开始维护的：for i in range(k//2-1, -1, -1)，i 就是我们要维护的堆的节点下标），遍历非叶子节点时要去维护非叶子节点所在堆中的位置，所以新定义一个函数initiate（这个函数肯定要传入当前节点，建好的堆结构，k）
• initiate（index，a，length）——初始化堆的函数。 index：维护当前堆的下标；a：数组（堆）；length：堆的节点个数。 由于可能要经过多次交换，所以设置一个temp变量去保存当前位置的值：temp = a[index]；然后就开始遍历【j = 2 * index+1：直接初始化为下标为index的节点的左节点】（循环的条件是 j 不超过堆的长度：while j < length）：，判断：if j+1 < length and a[j+1]>a[j]（该判断就是取出当前位置的左右孩子中节点值最大的节点，j是左孩子，j+1就是右孩子，若右孩子的值大于左孩子，就取右孩子的值：j+=1）；现已经找到了当前位置的左右孩子的最大值，就去比较：if a[j]>temp，就进行交换：a[index]=a[j]，index=j（更新 index 的值，index 就代表temp数字最终在堆中位置，当 j=j*2+1 执行后，index 和 j 的关系其实就是父节点和孩子节点的关系），else：break即可，由于我们是从下往上维护的就没有往下更新的必要了；退出for循环后，更新index所在位置的值即可：a[index]=temp
• 下面开始遍历n-k的部分：for i in range(k, n)，若当前位置的值小于大顶堆堆顶位置的值：if arr[i]<a[0]，a[0]=arr[i]，initiate(0,a,k)这样就将k个节点维护好了.维护好之后进行下一步，题目要求的结果是一个升序的操作，所以此刻的大顶堆还不是一个升序的，故下一步就是将大顶堆中的节点元素进行升序操作：分为两个过程：第一步交换，第二步固定（固定的操作其实是通过控制堆的节点个数去实现的）【每次都将堆顶元素与最后一个元素进行交换，交换之后的堆尾元素不参与堆的维护，再看剩下的部分】

三. 代码

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k > len(arr) or k == 0:
            return []
        n = len(arr)
        a = arr[:k]
        # 维护一个堆使其成为大顶堆
        for i in range(k//2-1, -1, -1):
            self.initiate(i, a, k)
        # 遍历剩下的n-k个数，判断是否要添加进堆中
        for i in range(k, n):
            if arr[i] < a[0]:
                a[0] = arr[i]
                self.initiate(0, a, k)
        # 将最后得到的大顶堆进行排序
        for i in range(len(a)-1, -1, -1):
            a[i], a[0] = a[0], a[i]
            self.initiate(0, a, i)
        return a

    def initiate(self, index, a, length):
        temp = a[index]
        j = index*2+1
        while j < length:
            if j+1 < length and a[j+1] > a[j]:
                j += 1
            if a[j] > temp:
                a[index] = a[j]
                index = j
            else:
                break
            j = j*2+1
        a[index] = temp

其中 initiate 函数也可写为：

def initiate(self, index, a, length):
        i, j = index, 2*index+1
        while j < length:
            if j+1 < length and a[j+1] > a[j]:
                j += 1
            if a[i] >= a[j]:
                break
            a[i], a[j] = a[j], a[i]
            i, j = j, 2*j+1