【C++】大根堆与小根堆

STL实现

C++中，大根堆和小根堆可以使用优先队列实现。

#include <queue>
priority_queue<int> pq1 // 大根堆
priority_queue<int, vector<int>, greater<int>> pq2 // 小根堆

该STL支持自定义比较函数，但与sort不同，不支持直接使用lambda函数。自定义的小根堆如下。

struct cmp {bool operator()(int a,int b) {return a > b;}}; // 自定义小根堆
priority_queue<int, vector<int>, cmp> pq;

更加复杂的自定义函数可以参照【LeetCode-1792】最大平均通过率。

手写

大小根堆的前备知识为堆排序，具体是使用堆排序中的节点上浮(swim)和下沉(sink)实现。整体的数据结构可以使用vector。push新的数值时，将其放在数组最后，然后与对应位置的父节点比较，进行上浮操作；pop堆顶数值时，将该节点与队尾数字交换，然后移除队尾数字，并对当前交换后的堆顶元素进行下沉操作。具体代码如下。

vector<int> pq{0}; // 首位填充一个数值，方便找根节点与子节点
void insert(int num) {
    pq.push_back(num);
    swim(pq.size() - 1);
}
int popTop() {
    swap(pq[1], pq.back());
    int res = pq.back();
    pq.pop_back();
    sink(1);
    return res;
}
int peekTop() {
    return pq[1];
}
bool isPrior(int i, int j) { // 自定义函数部分，此处实现的是大根堆
    return pq[i] > pq[j];
}
void swim(int i) { // 子节点与父节点比较
    while (i > 1 && isPrior(i, i / 2)) {
        swap(pq[i], pq[i / 2]);
        i /= 2;
    }
}
void sink(int i) { // 父节点与左右节点中最大的进行比较
    int n = pq.size();
    while (i * 2 < n) {
        int j = i * 2;
        if (j + 1 < n && isPrior(j + 1, j)) j++;
        if (isPrior(i , j)) break;
        swap(pq[i], pq[j]);
        i = j;
    }
}