实现 STL 中的优先队列

实现类似 STL 中的 priority_queue，包括 3 个操作：push, pop, top 。

• 用 vector 变长数组存放堆，索引从 0 开始。vec[i] 的左右子树分别为 vec[2 * i + 1], vec[2 * i + 2]；vec[i] 的父节点为 vec[(i - 1) / 2]
• 使用模版，支持大顶堆和小顶堆。

优先队列是用堆 (heap) 实现的数据结构，push, pop 的复杂度都是 \(O(\log{n})\)，top 复杂度是 \(O(1)\) 。

代码框架：

template <class T = int, class Comparator = less<T>> class PriorityQueue
{
  public:
    vector<T> vec;
    Comparator cmp;
  public:
    PriorityQueue() = default;

    T top() { return vec[0]; }
    void push(T x) { push_heap(vec, x); }
    T pop() { return pop_heap(vec); }
};

• cmp 是一个比较器，用于控制大顶堆还是小顶堆，同时也能实现自定义类型 T 的比较方法。
• cmp 默认是 less，并且默认情况下是大顶堆，与 STL 保持一致。
• vec 是存放堆的变长数组。

heapify

首先实现 heapify(vec, idx)，它的作用是将 vec[idx] 这一子树，调整为堆结构的子树。

void heapify(vector<T> &vec, int idx)
{
    int n = vec.size();
    // cur is the parent node, l and r are the children
    int l = 2 * idx + 1, r = 2 * idx + 2;
    int cur = idx;
    if (l < n && cmp(vec[cur], vec[l]))
        cur = l;
    if (r < n && cmp(vec[cur], vec[r]))
        cur = r;
    if (cur != idx)
    {
        std::swap(vec[idx], vec[cur]);
        heapify(vec, cur);
    }
}

如果需要将一个数组初始化为堆，那么需要：

void make_heap(vector<int> &vec)
{
    /* [n/2 + 1, n) is the leaf nodes. */
    for (int i = vec.size() / 2; i >= 0; --i)
        heapify(vec, i);
}

push_heap

push_heap(heap, x) 的作用是：

• heap 是一个 vector 存储的堆，x 是新元素。
• 将 x 压入 heap 中，并调整数组结构，使得新数组为一个合法堆。
• 执行结束，heap 的大小 (size) 增加 1 。

// Push element x into the `heap`, O(logn) time
void push_heap(vector<T> &heap, T x)
{
    int child, parent, n;
    heap.emplace_back(x);
    // We need to ajust the new element `heap[n - 1]` from bottom to up
    n = heap.size();
    child = n - 1, parent = (child - 1) / 2;
    while (child > 0 && cmp(heap[parent], heap[child]))
    {
        std::swap(heap[parent], heap[child]);
        child = parent;
        parent = (child - 1) / 2;
    }
}

pop_heap

pop_heap(vector &heap) 的作用是：

• 从 heap 中弹出并返回堆顶元素 heap[0]，然后调整剩余的元素，使得操作后也是一个合法堆。
• 执行结束后，heap 的大小减少 1 。

// Range [0, n - 1] is a heap, heap[0] is the top element.
T pop_heap(vector<T> &heap)
{
    assert(!heap.empty());
    std::swap(heap.front(), heap.back());
    T res = std::move(heap.back());
    heap.pop_back();
    // Re-adjust the heap in range of [0, size)
    if (heap.size() > 1)
        heapify(heap, 0);
    return res;
}

实现

有了上述 3 个操作，那么实现 PriorityQueue 是十分简单的。

template <class T = int, class Comparator = less<T>> class PriorityQueue
{
  public:
    vector<T> vec;
    Comparator cmp;
    PriorityQueue() = default;
    T top()
    {
        assert(!vec.empty());
        return vec[0];
    }
    void push(T x)
    {
        push_heap(vec, x);
    }
    T pop()
    {
        return pop_heap(vec);
    }
};

测试

template <class T1, class T2> bool cmp(T1 que, T2 expected)
{
    assert(que.vec.size() == expected.size());
    while (!expected.empty())
    {
        int a = que.top();
        int b = expected.top();
        assert(a == b);
        que.pop(), expected.pop();
    }
    return true;
}

template <class T1, class T2> void runTest(T1 &que, T2 &expected)
{
    int n = 1000;
    for (int i = 0; i < n; ++i)
    {
        int x = random();
        que.push(x);
        expected.push(x);
        assert(cmp(que, expected));
    }
}

int main()
{
    // Max-heap
    PriorityQueue<int> que;
    priority_queue<int> expected;
    // Min-heap
    PriorityQueue<int, greater<int>> que2;
    priority_queue<int, vector<int>, greater<int>> expected2;

    runTest(que, expected);
    runTest(que2, expected2);
}