python 中堆的使用

一、python 中堆的使用

1.1 heapq模块

相关文章：heapq — 堆队列算法
模块heapq中一些重要的函数
函数描述

• heappush(heap, x) 将x压入堆中 \n
• heappop(heap) 从堆中弹出最小的元素
• heapify(heap) 让列表具备堆特征
• heapreplace(heap, x) 弹出最小的元素，并将x压入堆中
• nlargest(n, iter) 返回iter中n个最大的元素
• nsmallest(n, iter) 返回iter中n个最小的元素

1.2 堆排序

堆排序适合关键字较多的情况（如n > 100）。
例如，在一亿个数中选出前100个最大值？
首先使用一个大小为100的数组，读入前100个数，建立小顶堆，而后依次读入余下的数，若小于堆顶则舍弃，否则用该数取代堆顶并重新调整堆，待数据读取完毕，堆中100个数即为所求。
堆排序建立时间复杂度O(nlogn)
大根堆：完全二叉树中，根>=左孩子，右孩子
小根堆：完全二叉树中，根<=左孩子，右孩子
对于一个大根堆（小根堆）来说，要符合
1、是一个完全二叉树
2、所有父节点均大于（小于）子节点
堆排序 可以通过将所有值推入堆中然后每次弹出一个最小值项来实现。 默认构建小顶堆

import heapq
# 实现堆排序
def heapsort(iterable):
     h = []
     for value in iterable:
         heapq.heappush(h, value)
     return [heapq.heappop(h) for i in range(len(h))]
heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])

堆元素可以为元组。 这适用于将比较值（例如任务优先级）与跟踪的主记录进行赋值的场合:

h = []
heapq.heappush(h,(5,'write code'))
heapq.heappush(h,(7,'write code'))
heapq.heappush(h,(1,'write spec'))
heapq.heappush(h,(3,'create tests'))
# 按元祖中的第一个比较
print(heapq.heappop(h)) 
print(heapq.heappop(h))

二、堆相关题目

347. 前 K 个高频元素

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

示例 1:

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

示例 2:

输入: nums = [1], k = 1
输出: [1]

AC如下

# 时间复杂度：O(nlogk)
# 空间复杂度：O(n)
import heapq
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
            整体思路：哈希 + 最小堆
        """
        # 使用map统计元素出现频率
        mp = {}
        for data in nums:
            mp[data] = mp.get(data, 0) + 1

        # 定义一个小顶堆，大小为k
        pri_que = [] 
        # 用固定大小为k的小顶堆，扫面所有频率的数值
        for key,value in mp.items():
            heapq.heappush(pri_que, (value,key))
             #如果堆的大小大于了K，则队列弹出，保证堆的大小一直为k
            if len(pri_que) > k: 
                heapq.heappop(pri_que)
        
        res = []
        for i in range(k):
            res.append(heapq.heappop(pri_que)[1])
        return res