堆(优先队列)是一种基于完全二叉树的动态数据结构,核心特性是快速获取最值(大根堆获取最大值,小根堆获取最小值),插入和删除操作的时间复杂度均为O(logn)O(\log n)O(logn)。它广泛应用于“动态维护最值”“Top-K 问题”“中位数维护”等场景,是处理动态数据的高效工具。本文通过4道经典题目,拆解堆在不同场景下的解题思路与代码实现。
一、最后一块石头的重量
题目描述:
有一堆石头,每回合选两块最重的石头粉碎:若重量相等则完全粉碎,否则剩下重量为两者差值的石头。返回最后剩下的石头重量(无剩余则返回0)。
示例:
- 输入:
stones = [2,7,4,1,8,1],输出:1(粉碎过程:8-7=1→4-2=2→2-1=1→1-1=0→剩1)
解题思路:
用大根堆维护石头重量,每次取最大的两块处理:
- 将所有石头重量入大根堆。
- 当堆中元素数>1时,取出最大的两块
a和b(a ≥ b):- 若
a > b,将a - b入堆; - 若
a == b,直接丢弃两块。
- 若
- 最终堆中若有元素则返回堆顶,否则返回0。
完整代码:
classSolution{public:intlastStoneWeight(vector<int>&stones){priority_queue<int>heap;// 大根堆(默认)for(autox:stones)heap.push(x);while(heap.size()>1){inta=heap.top();heap.pop();intb=heap.top();heap.pop();if(a>b)heap.push(a-b);}returnheap.size()?heap.top():0;}};复杂度分析:
- 时间复杂度:O(nlogn)O(n\log n)O(nlogn),
n为石头数量,每次入堆/出堆操作时间为O(logn)O(\log n)O(logn),最多执行nnn次。 - 空间复杂度:O(n)O(n)O(n),堆存储所有石头重量。
二、数据流中的第K大元素
题目描述:
设计一个类,动态维护数据流中的第K大元素(排序后的第K大,非第K个不同元素)。实现KthLargest类,包含初始化和添加元素后返回第K大的方法。
示例:
- 初始化:
k=3, nums=[4,5,8,2],添加3→返回4,添加5→返回5,添加10→返回5。
解题思路:
用小根堆维护“前K大的元素”,堆顶即为第K大元素:
- 初始化时,将所有元素入堆,若堆大小超过K则弹出堆顶(保留前K大的元素)。
- 添加元素时,将新元素入堆,若堆大小超过K则弹出堆顶,堆顶即为当前第K大元素。
完整代码:
classKthLargest{int_k;priority_queue<int,vector<int>,greater<int>>heap;// 小根堆public:KthLargest(intk,vector<int>&nums){_k=k;for(auto&x:nums){heap.push(x);if(heap.size()>_k)heap.pop();}}intadd(intval){heap.push(val);if(heap.size()>_k)heap.pop();returnheap.top();}};复杂度分析:
- 初始化时间:O(nlogK)O(n\log K)O(nlogK),
n为初始元素数,每个元素入堆/出堆时间为O(logK)O(\log K)O(logK)。 - 添加元素时间:O(logK)O(\log K)O(logK),每次入堆/出堆时间为O(logK)O(\log K)O(logK)。
- 空间复杂度:O(K)O(K)O(K),堆最多存储K个元素。
三、前K个高频单词
题目描述:
给定单词列表words和整数k,返回前K个出现次数最多的单词(频率相同按字典序升序排列)。
示例:
- 输入:
words = ["i","love","leetcode","i","love","coding"], k=2,输出:["i","love"](频率均为2,字典序i < love)
解题思路:
哈希表统计频率 + 小根堆维护前K个高频单词:
- 用哈希表统计每个单词的出现频率。
- 定义小根堆的比较规则:
- 频率不同时,频率小的优先出堆;
- 频率相同时,字典序大的优先出堆(保证堆顶是“频率最小/字典序最大”的候选,弹出后保留前K个)。
- 遍历哈希表,将“单词-频率”入堆,若堆大小超过K则弹出堆顶。
- 逆序收集堆中元素(因小根堆弹出的是较小的元素,需反转得到从大到小的顺序)。
完整代码:
classSolution{typedefpair<string,int>PSI;structcmp{booloperator()(constPSI a,constPSI b){if(a.second==b.second)returna.first<b.first;// 频率相同,字典序大的优先出堆elsereturna.second>b.second;// 频率小的优先出堆}};public:vector<string>topKFrequent(vector<string>&words,intk){unordered_map<string,int>hash;for(auto&x:words)hash[x]++;priority_queue<PSI,vector<PSI>,cmp>heap;for(auto&psi:hash){heap.push(psi);if(heap.size()>k)heap.pop();}vector<string>ret(k);for(inti=heap.size()-1;i>=0;i--){ret[i]=heap.top().first;heap.pop();}returnret;}};复杂度分析:
- 时间复杂度:O(mlogk)O(m\log k)O(mlogk),
m为不同单词的数量,每个单词入堆/出堆时间为O(logk)O(\log k)O(logk)。 - 空间复杂度:O(m+k)O(m + k)O(m+k),哈希表存储所有单词频率,堆存储K个单词。
四、数据流的中位数
题目描述:
设计一个类,动态维护数据流的中位数(奇数个元素取中间值,偶数个取中间两个的平均值)。实现MedianFinder类,包含添加元素和获取中位数的方法。
示例:
- 添加
1→添加2→中位数1.5→添加3→中位数2.0
解题思路:
用两个堆维护数据流的左右两部分:
- 大根堆
left:存储左半部分元素(≤中位数),堆顶为左半部分最大值; - 小根堆
right:存储右半部分元素(≥中位数),堆顶为右半部分最小值; - 保持平衡规则:
- 总元素数为偶数时,
left.size() == right.size(); - 总元素数为奇数时,
left.size() = right.size() + 1(中位数为left.top());
- 总元素数为偶数时,
- 添加元素时,根据元素与堆顶的大小关系选择入堆,并调整堆的大小以保持平衡。
完整代码:
classMedianFinder{priority_queue<int>left;// 大根堆(左半部分)priority_queue<int,vector<int>,greater<int>>right;// 小根堆(右半部分)public:MedianFinder(){}voidaddNum(intnum){if(left.size()==right.size()){if(left.empty()||num<left.top()){left.push(num);}else{right.push(num);left.push(right.top());right.pop();}}else{if(num<left.top()){left.push(num);right.push(left.top());left.pop();}else{right.push(num);}}}doublefindMedian(){if(left.size()==right.size())return(left.top()+right.top())/2.0;elsereturnleft.top();}};复杂度分析:
- 添加元素时间:O(logn)O(\log n)O(logn),每次入堆/出堆时间为O(logn)O(\log n)O(logn)。
- 获取中位数时间:O(1)O(1)O(1),直接取堆顶计算。
- 空间复杂度:O(n)O(n)O(n),两个堆存储所有元素。