【算法优选】优先级队列（堆）的经典应用

本文链接：https://blog.youkuaiyun.com/Java_ESS/article/details/148874276

掌握优先级队列的核心思想，高效解决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(log⁡K)O(\log K)$ ，get()操作 $O (1)$
空间复杂度： $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;
    }
}