ACM 学习篇 08：堆与优先队列

上一篇讲了并查集，它擅长处理连通性和集合合并。今天来看另一个非常实用的数据结构：堆（Heap），它是实现优先队列最常见的方式。

为什么需要堆

想象一个场景：不断有新任务进来，每个任务有不同的优先级，你每次要挑出当前优先级最高的任务来做。

如果用普通数组或链表，每次取最大值都要遍历一遍，时间复杂度是 O(n)。如果用排序后的数组，每次插入要移动元素，也很慢。

堆就是为了解决这类"动态维护最值"的问题而生的：它能在 O(log n) 的时间内插入元素和取出最值。

什么是二叉堆

二叉堆本质上是一棵完全二叉树，它用数组来存储，并且满足堆性质：

大顶堆：每个节点的值都大于等于它的左右子节点的值；
小顶堆：每个节点的值都小于等于它的左右子节点的值。

因为堆是完全二叉树，所以可以用数组紧凑地表示。如果根节点在下标 1 的位置，那么：

节点 i 的左子节点在 2 * i；
节点 i 的右子节点在 2 * i + 1；
节点 i 的父节点在 i / 2。

C++ 中的优先队列

C++ STL 提供了 priority_queue，默认是一个大顶堆：

priority_queue<int> maxHeap;  // 大顶堆
priority_queue<int, vector<int>, greater<int>> minHeap;  // 小顶堆

常用操作：

maxHeap.push(10);   // 插入元素，O(log n)
maxHeap.top();      // 取堆顶元素（最大值），O(1)
maxHeap.pop();      // 删除堆顶元素，O(log n)
maxHeap.size();     // 元素个数
maxHeap.empty();    // 是否为空

手写一个最小堆

虽然 STL 已经提供了优先队列，但理解堆的内部实现对你很有帮助。下面手写一个小顶堆：

class MinHeap {
private:
    vector<int> heap;

    void swim(int k) {  // 上浮操作
        while (k > 1 && heap[k] < heap[k / 2]) {
            swap(heap[k], heap[k / 2]);
            k /= 2;
        }
    }

    void sink(int k) {  // 下沉操作
        int n = heap.size() - 1;
        while (2 * k <= n) {
            int j = 2 * k;
            if (j < n && heap[j + 1] < heap[j]) j++;  // 选较小的子节点
            if (heap[k] <= heap[j]) break;
            swap(heap[k], heap[j]);
            k = j;
        }
    }

public:
    MinHeap() { heap.push_back(0); }  // 从下标1开始使用

    void push(int x) {
        heap.push_back(x);
        swim(heap.size() - 1);
    }

    int top() {
        return heap[1];
    }

    void pop() {
        int n = heap.size() - 1;
        swap(heap[1], heap[n]);
        heap.pop_back();
        sink(1);
    }

    bool empty() {
        return heap.size() == 1;
    }
};

核心思路只有两个：

上浮 swim：新元素放到末尾，然后一路向上与父节点比较，如果违反堆性质就交换；
下沉 sink：堆顶元素与末尾交换后弹出，然后从顶部一路向下，与较小的子节点比较交换。

典型题：数组中第 K 大的元素

给一个未排序的数组，找出其中第 K 大的元素。用小顶堆非常优雅：

int findKthLargest(vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<int>> minHeap;

    for (int x : nums) {
        minHeap.push(x);
        if (minHeap.size() > k) {
            minHeap.pop();  // 只保留最大的 k 个元素
        }
    }
    return minHeap.top();
}

思路：维护一个大小为 k 的小顶堆。堆顶就是当前已遍历元素中第 k 大的那个。遍历完后，堆顶就是答案。时间复杂度 O(n log k)。

典型题：合并 K 个有序链表

给 k 个有序链表，把它们合并成一条有序链表。这是堆的经典应用：

struct Compare {
    bool operator()(ListNode* a, ListNode* b) {
        return a->val > b->val;  // 小顶堆
    }
};

ListNode* mergeKLists(vector<ListNode*>& lists) {
    priority_queue<ListNode*, vector<ListNode*>, Compare> minHeap;

    for (auto head : lists) {
        if (head) minHeap.push(head);
    }

    ListNode dummy(0);
    ListNode* tail = &dummy;

    while (!minHeap.empty()) {
        ListNode* cur = minHeap.top();
        minHeap.pop();
        tail->next = cur;
        tail = cur;
        if (cur->next) minHeap.push(cur->next);
    }

    return dummy.next;
}

思路：每个链表先取一个头节点放入小顶堆，每次取出最小的接到结果里，然后把该节点的下一个节点重新入堆。总时间复杂度 O(N log k)，N 是所有节点总数。

典型题：数据流的中位数

不断有数据流入，每次要快速返回当前所有数据的中位数。用两个堆来做：

class MedianFinder {
private:
    priority_queue<int> maxHeap;           // 存较小的一半，堆顶是这一半的最大值
    priority_queue<int, vector<int>, greater<int>> minHeap;  // 存较大的一半，堆顶是这一半的最小值

public:
    void addNum(int num) {
        maxHeap.push(num);
        minHeap.push(maxHeap.top());
        maxHeap.pop();

        if (maxHeap.size() < minHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }

    double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.top();
        }
        return (maxHeap.top() + minHeap.top()) / 2.0;
    }
};

思路：大顶堆存较小的一半，小顶堆存较大的一半。保持大顶堆的元素个数等于或比小顶堆多一个，中位数就在两个堆顶之间。插入 O(log n)，查询 O(1)。

堆的其他应用

Top-K 问题：从海量数据中找最大的 K 个或最小的 K 个；
堆排序：先建堆，然后不断弹出堆顶，时间复杂度 O(n log n)；
Dijkstra 最短路径：用小顶堆选下一个距离最小的节点；
Prim 最小生成树：用小顶堆选下一条权值最小的边；
任务调度：按优先级取出任务执行。

堆的易错点

大顶堆还是小顶堆：不要搞反了，找 Top-K 大用小顶堆，找 Top-K 小用大顶堆；
自定义比较器：priority_queue 的比较器写法容易记反；
下标从 0 还是从 1 开始：手写堆时注意父节点和子节点的下标计算；
重复元素：堆允许重复元素，取出时要注意是否需要去重；
空间换时间：堆虽然速度快，但需要额外的数组空间。

什么时候用堆

看到这些关键词，就可以考虑堆：

动态维护最大值或最小值；
第 K 大、第 K 小、Top-K；
合并多个有序序列；
最短路径、最小生成树等图算法；
需要按优先级处理任务。

这一篇先记住什么

堆是一棵完全二叉树，用数组存储，满足堆性质；
两个核心操作：上浮 swim 和下沉 sink；
C++ 中 priority_queue 默认大顶堆，要小顶堆需指定 greater<T>；
堆的插入和删除都是 O(log n)，取堆顶是 O(1)；
经典应用：Top-K、合并 K 个有序链表、数据流中位数、Dijkstra。

下一篇继续数据结构：二叉搜索树（BST）。它能在 O(log n) 的平均时间内完成查找、插入和删除，同时保持元素有序，是很多更高级数据结构的基础。