堆与优先队列

概念讲解

什么是堆?

堆是一种完全二叉树,满足堆性质:

  • 大根堆: 父节点 >= 子节点
  • 小根堆: 父节点 <= 子节点

什么是优先队列?

优先队列是一种每次取出最大(或最小)元素的队列。

STL priority_queue

基本用法

123456789101112131415#include <queue>
using namespace std;

// 大根堆(默认)
priority_queue<int> pq;

// 小根堆
priority_queue<int, vector<int>, greater<int>> pq2;

// 操作
pq.push(1);    // 入队
pq.pop();      // 出队
pq.top();      // 获取堆顶
pq.empty();    // 判断是否为空
pq.size();     // 获取大小

自定义比较器

123456789101112struct Student {
    string name;
    int score;
};

struct cmp {
    bool operator()(Student a, Student b) {
        return a.score < b.score;  // 按分数从大到小
    }
};

priority_queue<Student, vector<Student>, cmp> pq;

堆的操作

插入

123456void heapifyUp(int i) {
    while (i > 0 && heap[i] > heap[(i - 1) / 2]) {
        swap(heap[i], heap[(i - 1) / 2]);
        i = (i - 1) / 2;
    }
}

删除堆顶

123456789101112131415161718void heapifyDown(int i) {
    int largest = i;
    int left = 2 * i + 1;
    int right = 2 * i + 2;
    
    if (left < heap.size() && heap[left] > heap[largest]) {
        largest = left;
    }
    
    if (right < heap.size() && heap[right] > heap[largest]) {
        largest = right;
    }
    
    if (largest != i) {
        swap(heap[i], heap[largest]);
        heapifyDown(largest);
    }
}

堆排序

123456789101112void heapSort(int a[], int n) {
    // 建堆
    for (int i = n / 2 - 1; i >= 0; i--) {
        heapifyDown(a, n, i);
    }
    
    // 排序
    for (int i = n - 1; i > 0; i--) {
        swap(a[0], a[i]);
        heapifyDown(a, i, 0);
    }
}

时间复杂度: O(n log n)

应用场景

1. 求第 k 大元素

123456789101112int findKthLargest(vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<int>> pq;
    
    for (int num : nums) {
        pq.push(num);
        if (pq.size() > k) {
            pq.pop();
        }
    }
    
    return pq.top();
}

2. 合并 k 个有序链表

12345678910111213141516171819202122ListNode* mergeKLists(vector<ListNode*>& lists) {
    priority_queue<pair<int, ListNode*>, vector<pair<int, ListNode*>>, greater<>> pq;
    
    for (ListNode* node : lists) {
        if (node) pq.push({node->val, node});
    }
    
    ListNode dummy;
    ListNode* curr = &dummy;
    
    while (!pq.empty()) {
        auto [val, node] = pq.top();
        pq.pop();
        
        curr->next = node;
        curr = curr->next;
        
        if (node->next) pq.push({node->next->val, node->next});
    }
    
    return dummy.next;
}

考点分析

初赛考查方式

  1. 选择题: 堆的特性
  2. 程序阅读: 理解堆操作
  3. 编程题: 堆的应用

常见题型

  • 求第 k 大元素
  • 合并有序序列
  • 任务调度

例题精讲

例题 1:求第 k 大元素

题目: 求数组中第 k 大的元素。

解答:

123456789101112131415161718192021222324252627#include <iostream>
#include <queue>
#include <vector>
using namespace std;

int main() {
    int n, k;
    cin >> n >> k;
    
    vector<int> nums(n);
    for (int i = 0; i < n; i++) {
        cin >> nums[i];
    }
    
    priority_queue<int, vector<int>, greater<int>> pq;
    
    for (int num : nums) {
        pq.push(num);
        if (pq.size() > k) {
            pq.pop();
        }
    }
    
    cout << pq.top() << endl;
    
    return 0;
}

练习题

  1. 洛谷 P3378 堆
查看答案

思路:使用 STL 优先队列实现小根堆。

123456789101112131415161718192021222324252627#include <iostream>
#include <queue>
using namespace std;

int main() {
    int n;
    cin >> n;
    
    priority_queue<int, vector<int>, greater<int>> pq;
    
    while (n--) {
        int op;
        cin >> op;
        
        if (op == 1) {
            int x;
            cin >> x;
            pq.push(x);
        } else if (op == 2) {
            cout << pq.top() << endl;
        } else {
            pq.pop();
        }
    }
    
    return 0;
}
  1. 洛谷 P1090 合并果子
查看答案

思路:每次合并最小的两堆,使用小根堆。

12345678910111213141516171819202122232425262728#include <iostream>
#include <queue>
using namespace std;

int main() {
    int n;
    cin >> n;
    
    priority_queue<int, vector<int>, greater<int>> pq;
    
    for (int i = 0; i < n; i++) {
        int x;
        cin >> x;
        pq.push(x);
    }
    
    int ans = 0;
    while (pq.size() > 1) {
        int a = pq.top(); pq.pop();
        int b = pq.top(); pq.pop();
        
        ans += a + b;
        pq.push(a + b);
    }
    
    cout << ans << endl;
    return 0;
}
  1. 洛谷 P1801 黑匣子
查看答案

思路:使用两个堆,大根堆存较小的一半,小根堆存较大的一半。

123456789101112131415161718192021222324252627282930313233343536373839#include <iostream>
#include <queue>
using namespace std;

int main() {
    int m, n;
    cin >> m >> n;
    
    vector<int> a(m);
    for (int i = 0; i < m; i++) cin >> a[i];
    
    vector<int> u(n);
    for (int i = 0; i < n; i++) cin >> u[i];
    
    priority_queue<int> maxHeap;  // 存较小的一半
    priority_queue<int, vector<int>, greater<int>> minHeap;  // 存较大的一半
    
    int idx = 0;
    for (int i = 0; i < m; i++) {
        maxHeap.push(a[i]);
        
        // 将 maxHeap 的最大值移到 minHeap
        minHeap.push(maxHeap.top());
        maxHeap.pop();
        
        // 保持 maxHeap 的大小 >= minHeap
        while (maxHeap.size() < minHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
        
        while (idx < n && u[idx] == i + 1) {
            cout << maxHeap.top() << endl;
            idx++;
        }
    }
    
    return 0;
}

易错点

⚠️ 注意:

  1. 大根堆和小根堆的区别
  2. 优先队列的 top() 不是 front()
  3. 堆排序是不稳定排序
  4. 建堆的时间复杂度是 O(n)