优先队列详解

news2025/6/13 19:23:50

优先队列是计算机科学中的一种抽象数据类型，它是一种队列：元素拥有优先级，优先级最高的元素最先得到服务；优先级相同的元素按照在集合中的顺序得到服务。优先队列有两种主要的实现方法：堆和二叉搜索树。

简单来说，优先队列是一种特殊的队列，不同于先进先出的普通队列，优先队列中的每个元素都有一定的优先级，优先级最高的元素最先出队。故而优先队列非常适合用来解决一些涉及元素优先级的问题。

在C++的STL中，priority_queue是实现优先队列的一个类模板。priority_queue中的元素默认按照大顶堆的方式进行排列，即出队顺序是优先级由高到低。

下面我们来介绍三个使用优先队列解决的典型问题：

1. 第K大元素查询

给定一个未排序的数组，找出第K大的元素。例如给定数组[3,2,1,5,6,4]和k = 2，返回5.

代码实现及注释:

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

int findKthLargest(vector<int>& nums, int k) {
    // 创建一个小顶堆，当堆的大小超过k时，弹出堆顶元素
    priority_queue<int, vector<int>, greater<int>> pq;
    for (int num : nums) {
        pq.push(num);
        if (pq.size() > k) {
            pq.pop();
        }
    }
    // 当所有元素都已入堆后，堆顶元素就是数组中第k大的元素
    return pq.top();
}

int main() {
    vector<int> nums = {3,2,1,5,6,4};
    int k = 2;
    cout << "The Kth largest number is: " << findKthLargest(nums, k) << endl;
    
    return 0;
}

2. 合并K个排序链表

给定K个已排序的链表，将它们合并为一个新的排序链表。

代码实现及注释:

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

struct ListNode {
    int val;
    ListNode *next;
    ListNode() : val(0), next(NULL) {}
};

struct compare {
    bool operator()(ListNode* l1, ListNode* l2) {
        return l1->val > l2->val;
    }
};

ListNode* mergeKLists(vector<ListNode*>& lists) {
    // 创建一个小顶堆，把所有链表的头节点放入堆中
    priority_queue<ListNode*, vector<ListNode*>, compare> pq;
    for (ListNode* list : lists) {
        if (list) {
            pq.push(list);
        }
    }

    ListNode *dummy = new ListNode();
    ListNode *cur = dummy;

    while (!pq.empty()) {
        // 取出当前最小的节点，存入新链表中
        ListNode* node = pq.top(); 
        pq.pop();
        cur->next = node;
        cur = node;

        if (node->next) {
            pq.push(node->next); // 将当前最小节点的后续节点放入优先队列中
        }
    }

    return dummy->next; // 返回去掉哑节点头的链表
}

3. 数据流的中位数

给定一个数据流，寻找数据流的中位数。例如，数据流[2,3,4]的中位数是3，数据流[2,3]的中位数是2.5.

代码实现及注释:

#include <queue>
using namespace std;

class MedianFinder {
public:
    /** initialize your data structure here. */
    MedianFinder() {

    }

    void addNum(int num) {
        left.push(num);
        right.push(-left.top());
        left.pop();

        if (left.size() < right.size()) {
            left.push(-right.top());
            right.pop();
        }
    }

    double findMedian() {
        if (left.size() > right.size()) {
            return left.top();
        } else {
            return (left.top() - right.top()) / 2.0;
        }
    }

private:
    // 左侧大顶堆，右侧小顶堆，保证大顶堆的所有元素都小于小顶堆，且大顶堆的元素可以多一个
    priority_queue<int> left;
    priority_queue<int> right;
};