【专题讲解】手撕堆

引入问题是leetcode347题，前K个高频元素
说一千到一万还是一个堆问题。虽然python有堆函数可以直接调用。但是我们还是学习一下利用堆实现优先队列的方法。

堆性质：

• 堆必须是完全二叉树(保证利用数组存储堆时，由下标i快速找到父节点(i−1)/2，左右子节点2∗i+1,2∗i+2。（从0开始存储）

• 堆中每个节点必须大于等于(小于等于)子节点

• 支持插入、删除堆顶元素(并且要堆化)

• 如果数组存储n个元素，最后一个非叶子节点下标为n/2−1

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        # hashmap 统计频率
        freq_count = {}
        for num in nums:
            if num in freq_count:
                freq_count[num] += 1
            else:
                freq_count[num] = 1

        def sift_up(arr, k):
            """ 时间复杂度 O(logk) k 为堆的规模"""
            new_index, new_val = k-1, arr[k-1]
            while (new_index > 0 and arr[(new_index-1)//2][1] > new_val[1]):
                arr[new_index] = arr[(new_index-1)//2]
                new_index = (new_index-1)//2
            arr[new_index] = new_val # 这里采用的是类似插入排序的赋值交换

        def sift_down(arr, root, k):
            """ O(logk). 左节点index 2*root+1, 右节点 2*root+2, 父节点 (child-1)//2"""
            root_val = arr[root]
            while (2*root+1 < k):
                child = 2 * root + 1
                # 小顶锥 用 >，大顶锥 用 <
                if child+1 < k and arr[child][1] > arr[child+1][1]:
                    child += 1
                ## 两个子节点中小与根节点比较
                if root_val[1] > arr[child][1]:
                    arr[root] = arr[child]
                    root = child # 继续向下检查
                else: 
                    break # 如果到这里没乱序，不用再检查后续子节点
            ## 新数值下沉到该位置
            arr[root] = root_val

        # 注意构造规模为k的堆, 时间复杂度O(n)，因为堆的规模是从0开始增长的
        freq_list = list(freq_count.items())
        # print(freq_list) [(1, 3), (2, 2), (3, 1)]第一项是key,第二项freq
        ## 采用小顶堆，确定了kk个以后，只要又大于堆顶的入堆，堆顶移除
        min_heap = []
        for i in range(k):
            min_heap.append(freq_list[i])
            sift_up(min_heap, i+1)

        # 遍历剩下元素，大于堆顶入堆，下沉维护小顶堆
        for item in freq_list[k:]:
            priority = item[1]
            if priority > min_heap[0][1]:
                min_heap[0] = item
                sift_down(min_heap, 0, k)

        return [item[0] for item in min_heap]