掌握优先级队列的核心思想,高效解决TopK/中位数等高频算法问题!
一、最后一块石头的重量
问题描述
每次选择最重的两块石头碰撞,若重量相同则消失,不同则留下差值。求最终剩余石头的重量。
示例:
输入:[2,7,4,1,8,1]
→ 碰撞过程:
- 8 vs 7 → 余1 →
[2,4,1,1,1]
- 4 vs 2 → 余2 →
[1,1,1,2]
- 2 vs 1 → 余1 →
[1,1,1]
- 1 vs 1 → 消失 → 最终结果:
1
解题思路
- 大顶堆存储:将所有石头放入大顶堆(Java的
PriorityQueue
默认小顶堆,需重写比较器) - 循环碰撞:每次弹出堆顶两个元素进行碰撞
- 差值入堆:若碰撞后剩余重量>0,将差值重新入堆
- 终止条件:堆中元素≤1时停止
代码实现
class Solution {
public int lastStoneWeight(int[] stones) {
PriorityQueue<Integer> heap = new PriorityQueue<>((a, b) -> b - a); // 大顶堆
for (int stone : stones) heap.offer(stone);
while (heap.size() > 1) {
int y = heap.poll(); // 最重石头
int x = heap.poll(); // 次重石头
if (y > x) heap.offer(y - x); // 碰撞后剩余
}
return heap.isEmpty() ? 0 : heap.peek();
}
}
二、数据流中的第 K 大元素
问题特点
- 动态数据流:持续添加新元素
- 实时获取当前第K大的元素
解决方案:小顶堆维护TopK
- 初始化:创建大小为K的小顶堆
- 添加元素:
- 堆未满时直接入堆
- 堆满后,若新元素>堆顶则替换堆顶
- 获取结果:堆顶即第K大元素
复杂度分析
- 时间复杂度:
add()
操作O(logK)O(\log K)O(logK),get()
操作O(1)O(1)O(1) - 空间复杂度:O(K)O(K)O(K)
class KthLargest {
private final PriorityQueue<Integer> minHeap;
private final int k;
public KthLargest(int k, int[] nums) {
this.k = k;
minHeap = new PriorityQueue<>();
for (int num : nums) add(num); // 复用添加逻辑
}
public int add(int val) {
if (minHeap.size() < k) {
minHeap.offer(val);
} else if (val > minHeap.peek()) {
minHeap.poll(); // 移除堆顶
minHeap.offer(val); // 加入新元素
}
return minHeap.peek();
}
}
三、前K个高频单词
双关键排序挑战
- 频率降序:出现次数多的在前
- 字典序升序:频率相同时按字母顺序排序
解题步骤
- 统计频率:HashMap记录单词出现次数
- 堆排序:
- 小顶堆维护TopK(按频率升序/字典序降序)
- 自定义比较器实现双条件排序
- 结果反转:堆中元素弹出后需反转顺序
核心比较器
Comparator<String> comp = (a, b) ->
map.get(a).equals(map.get(b)) ?
b.compareTo(a) : // 频率相同:按字典序降序(Z->A)
map.get(a) - map.get(b); // 频率升序
完整代码
class Solution {
public List<String> topKFrequent(String[] words, int k) {
Map<String, Integer> map = new HashMap<>();
for (String word : words)
map.put(word, map.getOrDefault(word, 0) + 1);
// 小顶堆:频率升序,字典序降序
PriorityQueue<String> heap = new PriorityQueue<>(
(a, b) -> map.get(a).equals(map.get(b)) ?
b.compareTo(a) : map.get(a) - map.get(b)
);
for (String word : map.keySet()) {
heap.offer(word);
if (heap.size() > k) heap.poll(); // 维护堆大小
}
// 反转结果
LinkedList<String> res = new LinkedList<>();
while (!heap.isEmpty()) res.addFirst(heap.poll());
return res;
}
}
四、数据流的中位数
核心难点
动态数据流中实时获取中位数,需同时支持:
- 快速插入:添加新元素
- 快速查询:获取当前中位数
双堆平衡策略
数据结构 | 存储内容 | 特点 |
---|---|---|
大顶堆 | 较小的一半数据 | 堆顶是左半段最大值 |
小顶堆 | 较大的一半数据 | 堆顶是右半段最小值 |
操作规则
-
添加元素:
-
获取中位数:
- 总数为奇:大顶堆堆顶
- 总数为偶:两堆堆顶平均值
代码实现
class MedianFinder {
private final PriorityQueue<Integer> maxHeap; // 存储较小值
private final PriorityQueue<Integer> minHeap; // 存储较大值
public MedianFinder() {
maxHeap = new PriorityQueue<>((a, b) -> b - a);
minHeap = new PriorityQueue<>();
}
public void addNum(int num) {
// 选择插入堆
if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
maxHeap.offer(num);
} else {
minHeap.offer(num);
}
// 平衡堆大小
if (maxHeap.size() > minHeap.size() + 1) {
minHeap.offer(maxHeap.poll());
} else if (minHeap.size() > maxHeap.size()) {
maxHeap.offer(minHeap.poll());
}
}
public double findMedian() {
if (maxHeap.size() > minHeap.size()) {
return maxHeap.peek();
}
return (maxHeap.peek() + minHeap.peek()) / 2.0;
}
}
复杂度分析
操作 | 时间复杂度 | 空间复杂度 |
---|---|---|
添加元素 | O(logN)O(\log N)O(logN) | O(N)O(N)O(N) |
查询中位数 | O(1)O(1)O(1) | - |
堆的本质是部分有序的结构化思维,在动态数据场景中,它能以最低成本维护关键有序信息。掌握这四类经典问题,即可解决90%优先级队列相关算法题!